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

lmcache/cli/commands/query/request.py

@@ -0,0 +1,357 @@
+# SPDX-License-Identifier: Apache-2.0
+"""HTTP request for VLLM Serving Engine."""
+
+# Standard
+from typing import Any, Optional
+import json
+import sys
+import time
+import urllib.error
+import urllib.request
+
+_MAX_ERR = 65536
+MetricValue = tuple[str, Any]
+MetricMap = dict[str, MetricValue]
+_METRIC_NAMES = {
+    "prompt_tokens": "Input tokens",
+    "output_tokens": "Output tokens",
+    "ttft_ms": "TTFT (ms)",
+    "tpot_ms_per_token": "TPOT (ms/token)",
+    "total_latency_ms": "Total latency (ms)",
+    "throughput_tokens_per_s": "Throughput (tokens/s)",
+    "model": "Model",
+}
+
+
+def _clip(text: str, limit: int = _MAX_ERR) -> str:
+    return (
+        text
+        if len(text) <= limit
+        else text[: max(0, limit - 24)] + "\n...(message truncated)..."
+    )
+
+
+def _info(msg: str) -> None:
+    print(f"lmcache query: {msg}", file=sys.stderr)
+
+
+def _openai_error(obj: dict[str, Any]) -> Optional[str]:
+    err = obj.get("error")
+    if err is None:
+        return None
+    if isinstance(err, str):
+        return err.strip() or None
+    if not isinstance(err, dict):
+        return _clip(str(err))
+    for key in ("message", "detail"):
+        val = err.get(key)
+        if not isinstance(val, str) or not val.strip():
+            continue
+        typ = err.get("type") or err.get("code")
+        if key == "message" and isinstance(typ, str) and typ.strip():
+            return f"{typ.strip()}: {val.strip()}"
+        return val.strip()
+    try:
+        return _clip(json.dumps(err, ensure_ascii=False))
+    except Exception:
+        return _clip(str(err))
+
+
+def _raise_openai_error(obj: dict[str, Any]) -> None:
+    msg = _openai_error(obj)
+    if msg:
+        raise RuntimeError(_clip(msg))
+
+
+def _raise_json_blob_error(blob: str) -> None:
+    s = blob.strip()
+    if not s.startswith("{"):
+        return
+    try:
+        obj = json.loads(s)
+    except json.JSONDecodeError:
+        return
+    if isinstance(obj, dict):
+        _raise_openai_error(obj)
+
+
+def _api_url(base: str, path: str) -> str:
+    base = base.strip()
+    if "://" not in base:
+        base = f"http://{base}"
+    base = base.rstrip("/")
+    return f"{base if base.endswith('/v1') else base + '/v1'}/{path}"
+
+
+def _read_json(url: str, timeout: float) -> dict[str, Any]:
+    try:
+        with urllib.request.urlopen(
+            urllib.request.Request(url, method="GET"), timeout=max(timeout + 2.0, 5.0)
+        ) as resp:
+            raw = resp.read().decode("utf-8", errors="replace")
+    except urllib.error.HTTPError as e:
+        body = e.read().decode("utf-8", errors="replace")[:512]
+        raise RuntimeError(
+            f"GET {url} failed (HTTP {e.code}): {body or 'no body'}"
+        ) from e
+    except urllib.error.URLError as e:
+        raise RuntimeError(f"GET {url} failed: {getattr(e, 'reason', e)}") from e
+    try:
+        obj = json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise RuntimeError(f"Invalid JSON from GET {url}: {e}") from e
+    if not isinstance(obj, dict):
+        raise RuntimeError(f"GET {url}: expected a JSON object")
+    return obj
+
+
+def _sse_piece(obj: dict[str, Any], chat: bool) -> str:
+    choices = obj.get("choices") or []
+    if not choices:
+        return ""
+    c0 = choices[0]
+    return (
+        str((c0.get("delta") or {}).get("content") or "")
+        if chat
+        else str(c0.get("text") or "")
+    )
+
+
+def _trim_misc_buffer(misc: list[str], limit: int = _MAX_ERR) -> None:
+    while misc and sum(map(len, misc)) > limit:
+        misc.pop(0)
+
+
+def _stream(
+    url: str,
+    body: dict[str, Any],
+    timeout: float,
+    *,
+    chat: bool,
+    max_tokens: int,
+) -> dict[str, Any]:
+    """POST with ``stream: true``; parse SSE; return TTFT/TPOT and token metrics."""
+    payload = {
+        **body,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+    }
+    req = urllib.request.Request(
+        url,
+        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
+        method="POST",
+        headers={"Content-Type": "application/json"},
+    )
+    t0, first_token_t, pieces, usage, misc = time.time(), None, [], None, []
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            while True:
+                raw = resp.readline()
+                if not raw:
+                    break
+                line = raw.decode("utf-8", errors="replace").strip()
+                if not line:
+                    continue
+                if not line.startswith("data:"):
+                    misc.append(line)
+                    _trim_misc_buffer(misc)
+                    continue
+                chunk = line[5:].strip()
+                if chunk == "[DONE]":
+                    break
+                try:
+                    obj = json.loads(chunk)
+                except json.JSONDecodeError:
+                    misc.append(chunk)
+                    _trim_misc_buffer(misc)
+                    continue
+                if not isinstance(obj, dict):
+                    continue
+                _raise_openai_error(obj)
+                piece = _sse_piece(obj, chat)
+                if piece:
+                    first_token_t = first_token_t or time.time()
+                    pieces.append(piece)
+
+                u_chunk = obj.get("usage")
+                if u_chunk is not None:
+                    usage = u_chunk
+            t1 = time.time()
+    except urllib.error.HTTPError as e:
+        err_body = e.read().decode("utf-8", errors="replace")
+        _raise_json_blob_error(err_body)
+        raise RuntimeError(
+            _clip(f"POST {url} failed (HTTP {e.code}):\n{_clip(err_body)}")
+        ) from e
+    except urllib.error.URLError as e:
+        raise RuntimeError(f"POST {url} failed: {getattr(e, 'reason', e)}") from e
+
+    misc_text = "\n".join(misc).strip()
+    _raise_json_blob_error(misc_text)
+    joined = "".join(pieces)
+    if not joined and usage is None:
+        raise RuntimeError(
+            _clip(f"No completion output from engine. Captured response:\n{misc_text}")
+            if misc_text
+            else "Empty response from engine (no SSE chunks parsed)."
+        )
+
+    u = usage or {}
+    prompt_tokens = int(u.get("prompt_tokens") or 0)
+    num_completion = int(u.get("completion_tokens") or 0)
+    # Match V2RequestSender: server count if present, else max_tokens cap.
+    num_generated = num_completion if num_completion > 0 else max_tokens
+    if first_token_t is None:
+        # Use total round-trip as a conservative TTFT approximation.
+        ttft_s = t1 - t0
+        decode_time = 0.0
+    else:
+        ttft_s = first_token_t - t0
+        decode_time = t1 - first_token_t
+    dt = t1 - t0
+    decoding_speed = (num_generated / decode_time) if decode_time > 0 else 0.0
+    tpot_s = (
+        (decode_time / num_generated) if num_generated > 0 and decode_time > 0 else 0.0
+    )
+    return {
+        "prompt_tokens": prompt_tokens,
+        "output_tokens": num_generated,
+        "ttft_ms": ttft_s * 1000.0,
+        "tpot_ms_per_token": tpot_s * 1000.0,
+        "total_latency_ms": dt * 1000.0,
+        "throughput_tokens_per_s": decoding_speed,
+    }
+
+
+def _missing_chat_template(exc: BaseException) -> bool:
+    msg = str(exc).lower()
+    return any(
+        s in msg
+        for s in (
+            "chat template",
+            "chat_template",
+            "chattemplate",
+            "template resolution",
+            "must provide a chat template",
+            "default chat template is no longer allowed",
+        )
+    )
+
+
+def _weak_completions_error(msg: str) -> bool:
+    msg = msg.lower()
+    return any(
+        s in msg
+        for s in (
+            "empty response from engine",
+            "no completion output from engine",
+            "no sse chunks parsed",
+        )
+    )
+
+
+class Request:
+    """Build and send one query request against an OpenAI-compatible endpoint."""
+
+    def __init__(
+        self,
+        base: str,
+        model: Optional[str],
+        max_tokens: int,
+        timeout: float,
+        *,
+        completions_only: bool = False,
+        chat_first: bool = False,
+    ) -> None:
+        self._base = base
+        self._model = model
+        self._max_tokens = max_tokens
+        self._timeout = timeout
+        self._completions_only = completions_only
+        self._chat_first = chat_first
+
+    def build_request(self, prompt: str) -> dict[str, Any]:
+        """Build request payload and metadata for the provided prompt."""
+        model = self._model or self._first_model_id()
+        return {
+            "base": self._base,
+            "model": model,
+            "prompt": prompt,
+            "max_tokens": self._max_tokens,
+            "timeout": self._timeout,
+            "completions_only": self._completions_only,
+            "chat_first": self._chat_first,
+        }
+
+    def send_request(self, prompt: str) -> dict[str, Any] | MetricMap:
+        """Send request and return stats."""
+        request_data = self.build_request(prompt)
+        stats = {
+            "model": request_data["model"],
+            **self._query_with_fallback(request_data),
+        }
+        return {key: (_METRIC_NAMES.get(key, key), stats[key]) for key in stats}
+
+    def _first_model_id(self) -> str:
+        """Return the first model ID from ``GET /v1/models``."""
+        obj = _read_json(_api_url(self._base, "models"), self._timeout)
+        data = obj.get("data")
+        if not isinstance(data, list) or not data:
+            raise RuntimeError(
+                "GET /v1/models returned no models; pass --model explicitly."
+            )
+        first = data[0]
+        if not isinstance(first, dict) or "id" not in first:
+            raise RuntimeError("GET /v1/models: first entry missing 'id'.")
+        return str(first["id"])
+
+    def _query_with_fallback(self, request_data: dict[str, Any]) -> dict[str, Any]:
+        """Send one query and fallback between completions/chat endpoints."""
+        if request_data["completions_only"]:
+            return self._query(request_data, chat=False)
+        try:
+            return self._query(request_data, chat=request_data["chat_first"])
+        except RuntimeError as first_err:
+            if request_data["chat_first"]:
+                if not _missing_chat_template(first_err):
+                    raise
+                _info(
+                    "chat API failed (no chat template); retrying with /v1/completions"
+                )
+                return self._query(request_data, chat=False)
+            _info("/v1/completions failed; retrying with /v1/chat/completions")
+            try:
+                return self._query(request_data, chat=True)
+            except RuntimeError as second_err:
+                if _weak_completions_error(str(first_err)) and _missing_chat_template(
+                    second_err
+                ):
+                    _info(
+                        "base / completion-only models: try `--completions` or "
+                        "an instruct model with a chat template."
+                    )
+                    raise second_err
+                raise RuntimeError(f"{first_err}; then {second_err}") from second_err
+
+    def _query(self, request_data: dict[str, Any], *, chat: bool) -> dict[str, Any]:
+        path = "chat/completions" if chat else "completions"
+        model = request_data["model"]
+        prompt = request_data["prompt"]
+        max_tokens = request_data["max_tokens"]
+        timeout = request_data["timeout"]
+        body = (
+            {
+                "model": model,
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": max_tokens,
+            }
+            if chat
+            else {"model": model, "prompt": prompt, "max_tokens": max_tokens}
+        )
+        return _stream(
+            _api_url(request_data["base"], path),
+            body,
+            timeout,
+            chat=chat,
+            max_tokens=max_tokens,
+        )

lmcache/cli/commands/server.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache server`` — launch the LMCache server (ZMQ + HTTP)."""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+
+
+class ServerCommand(BaseCommand):
+    """CLI command that launches the LMCache server (ZMQ + HTTP)."""
+
+    def name(self) -> str:
+        """Return the subcommand name.
+
+        Returns:
+            The string ``"server"``.
+        """
+        return "server"
+
+    def help(self) -> str:
+        """Return short help text.
+
+        Returns:
+            Help string shown by ``lmcache -h``.
+        """
+        return "Launch the LMCache server (ZMQ + HTTP)."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        """Add server-specific arguments to the parser.
+
+        Composes argument groups from the multiprocess, storage manager,
+        HTTP frontend, Prometheus, and telemetry config modules.
+        Silently skips argument registration when server dependencies
+        (e.g. CUDA extensions) are not installed; ``execute`` will then
+        print an actionable error.
+
+        Args:
+            parser: The ``ArgumentParser`` for this subcommand.
+        """
+        try:
+            # First Party
+            from lmcache.v1.distributed.config import add_storage_manager_args
+            from lmcache.v1.mp_observability.config import add_observability_args
+            from lmcache.v1.multiprocess.config import (
+                add_http_frontend_args,
+                add_mp_server_args,
+            )
+        except ImportError as e:
+            print(
+                f"Failed to import server dependencies: {e}. "
+                "Install the full lmcache package to use 'lmcache server'."
+            )
+            return
+
+        add_mp_server_args(parser)
+        add_storage_manager_args(parser)
+        add_http_frontend_args(parser)
+        add_observability_args(parser)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        """Parse CLI arguments into config objects and launch the HTTP server.
+
+        Args:
+            args: Parsed CLI arguments.
+
+        Raises:
+            SystemExit: When server dependencies are not installed.
+        """
+        # Standard
+        import sys
+
+        try:
+            # First Party
+            from lmcache.v1.distributed.config import parse_args_to_config
+            from lmcache.v1.mp_observability.config import (
+                parse_args_to_observability_config,
+            )
+            from lmcache.v1.multiprocess.config import (
+                parse_args_to_http_frontend_config,
+                parse_args_to_mp_server_config,
+            )
+            from lmcache.v1.multiprocess.http_server import run_http_server
+        except ImportError:
+            print(
+                "The 'lmcache server' command requires the full lmcache "
+                "installation with CUDA extensions.\n"
+                "Install with: pip install lmcache",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+        run_http_server(
+            http_config=parse_args_to_http_frontend_config(args),
+            mp_config=parse_args_to_mp_server_config(args),
+            storage_manager_config=parse_args_to_config(args),
+            obs_config=parse_args_to_observability_config(args),
+        )

lmcache/cli/commands/tool/__init__.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache tool`` command — offline analysis utilities.
+
+To add a new tool:
+
+1. Create ``lmcache/cli/commands/tool/<your_tool>.py`` with a ``register``
+   function that adds a sub-subcommand to the supplied ``subparsers``.
+2. Import it here and call ``your_tool.register(inner)`` inside
+   :meth:`ToolCommand.register`.
+"""
+
+# Standard
+import argparse
+import sys
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+from lmcache.cli.commands.tool import cache_simulator
+
+
+class ToolCommand(BaseCommand):
+    """CLI command for offline analysis tools bundled with LMCache."""
+
+    def name(self) -> str:
+        """Return the subcommand name."""
+        return "tool"
+
+    def help(self) -> str:
+        """Return short help text shown by ``lmcache -h``."""
+        return "Run offline analysis tools."
+
+    def add_arguments(self, _parser: argparse.ArgumentParser) -> None:
+        """No top-level arguments; all args are registered in register()."""
+
+    def register(self, subparsers: argparse._SubParsersAction) -> None:
+        """Register ``lmcache tool`` and all tool sub-subcommands.
+
+        Args:
+            subparsers: The subparsers action from the root parser.
+        """
+        parser = subparsers.add_parser(
+            self.name(),
+            help=self.help(),
+            description="Run offline analysis tools bundled with LMCache.",
+        )
+        inner = parser.add_subparsers(
+            dest="tool_name",
+            required=True,
+            metavar="{cache-simulator}",
+        )
+        cache_simulator.register(inner)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        """Dispatch is handled per-tool via parser.set_defaults(func=...).
+
+        This method is never called directly; each tool's register() binds
+        its own execute function as the dispatch target.
+
+        Args:
+            args: Parsed CLI arguments.
+        """
+        print(f"Unknown tool: {getattr(args, 'tool_name', '?')}", file=sys.stderr)
+        sys.exit(1)

lmcache/cli/commands/tool/cache_simulator.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache tool cache-simulator`` sub-subcommand wiring.
+
+Registers the ``simulate``, ``sweep``, and ``gen-dataset`` actions under
+``lmcache tool cache-simulator``.  Flag definitions and execution logic
+live entirely in the simulator modules:
+
+* :func:`~lmcache.tools.cache_simulator.simulator.add_simulate_arguments`
+* :func:`~lmcache.tools.cache_simulator.simulator.run_simulate`
+* :func:`~lmcache.tools.cache_simulator.plot_hit_rate.add_sweep_arguments`
+* :func:`~lmcache.tools.cache_simulator.plot_hit_rate.run_sweep`
+* :func:`~lmcache.tools.cache_simulator.gen_bench_dataset.add_gen_dataset_arguments`
+* :func:`~lmcache.tools.cache_simulator.gen_bench_dataset.run_gen_dataset`
+
+To add a new action, add ``add_*_arguments`` / ``run_*`` functions to the
+appropriate simulator module and wire them up in :func:`register`.
+"""
+
+# Standard
+import argparse
+
+
+def register(subparsers: argparse._SubParsersAction) -> None:
+    """Register ``cache-simulator`` and its ``simulate``/``sweep`` actions.
+
+    Imports are lazy so that matplotlib is not loaded at CLI startup.
+
+    Args:
+        subparsers: The subparsers action from the ``lmcache tool`` parser.
+    """
+    # Lazy imports — keeps CLI startup fast and makes cache-simulator optional.
+    # Missing sortedcontainers/matplotlib (plot extra) skips the subcommand.
+    try:
+        # First Party
+        from lmcache.tools.cache_simulator.gen_bench_dataset import (
+            add_gen_dataset_arguments,
+        )
+        from lmcache.tools.cache_simulator.plot_hit_rate import add_sweep_arguments
+        from lmcache.tools.cache_simulator.simulator import add_simulate_arguments
+    except ImportError:
+        return
+
+    cs_parser = subparsers.add_parser(
+        "cache-simulator",
+        help="Simulate KV-cache token hit rate from lookup-hash JSONL logs.",
+        description=(
+            "Replay LMCache lookup-hash JSONL logs through an LRU cache "
+            "to measure token hit rate."
+        ),
+    )
+    cs_sub = cs_parser.add_subparsers(
+        dest="cs_action",
+        required=True,
+        metavar="{simulate,sweep,gen-dataset}",
+    )
+
+    sim_parser = cs_sub.add_parser(
+        "simulate",
+        help=(
+            "Replay logs at a fixed cache capacity; print a text report "
+            "and save a 7-panel statistics PNG."
+        ),
+    )
+    add_simulate_arguments(sim_parser)
+    sim_parser.set_defaults(func=execute)
+
+    sweep_parser = cs_sub.add_parser(
+        "sweep",
+        help=(
+            "Sweep across a range of cache capacities and save a "
+            "hit-rate vs capacity PNG."
+        ),
+    )
+    add_sweep_arguments(sweep_parser)
+    sweep_parser.set_defaults(func=execute)
+
+    gen_parser = cs_sub.add_parser(
+        "gen-dataset",
+        help=(
+            "Generate a vllm bench serve custom dataset (JSONL) from "
+            "lookup-hash JSONL logs, preserving prefix-sharing structure."
+        ),
+    )
+    add_gen_dataset_arguments(gen_parser)
+    gen_parser.set_defaults(func=execute)
+
+
+def execute(args: argparse.Namespace) -> None:
+    """Dispatch to the correct cache-simulator action.
+
+    Args:
+        args: Parsed CLI arguments (includes ``cs_action``).
+    """
+    # Standard
+    import sys
+
+    # First Party
+    # Lazy imports — keeps CLI startup fast (avoids loading matplotlib)
+    from lmcache.tools.cache_simulator.plot_hit_rate import run_sweep
+    from lmcache.tools.cache_simulator.simulator import run_simulate
+
+    if args.cs_action == "simulate":
+        run_simulate(args)
+    elif args.cs_action == "sweep":
+        run_sweep(args)
+    elif args.cs_action == "gen-dataset":
+        # First Party
+        from lmcache.tools.cache_simulator.gen_bench_dataset import run_gen_dataset
+
+        run_gen_dataset(args)
+    else:
+        print(f"Unknown cache-simulator action: {args.cs_action}", file=sys.stderr)
+        sys.exit(1)