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

lmcache/cli/commands/ping.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``lmcache ping`` — liveness check for LMCache and vLLM servers.
+
+Usage::
+
+    lmcache ping kvcache --url http://localhost:8080
+    lmcache ping engine  --url http://localhost:8000
+"""
+
+# Standard
+import argparse
+import sys
+import time
+import urllib.error
+import urllib.request
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+from lmcache.cli.commands.describe import normalize_url
+
+# -------------------------------------------------------------------
+# Constants
+# -------------------------------------------------------------------
+
+HEALTH_ENDPOINTS: dict[str, str] = {
+    "kvcache": "/api/healthcheck",
+    "engine": "/health",
+}
+
+DEFAULT_URLS: dict[str, str] = {
+    "kvcache": "http://localhost:8080",
+    "engine": "http://localhost:8000",
+}
+
+TITLES: dict[str, str] = {
+    "kvcache": "Ping KV Cache",
+    "engine": "Ping Engine",
+}
+
+# -------------------------------------------------------------------
+# Helper
+# -------------------------------------------------------------------
+
+
+def ping(url: str, timeout: int = 10) -> tuple[str, float, str | None]:
+    """GET *url* and return ``(status, rtt_ms, error_msg)``.
+
+    Returns:
+        ``("OK", rtt_ms, None)`` on HTTP 200.
+        ``("FAIL", rtt_ms, detail)`` on any error.
+    """
+    start = time.monotonic()
+    try:
+        with urllib.request.urlopen(
+            urllib.request.Request(url), timeout=timeout
+        ) as resp:
+            rtt_ms = (time.monotonic() - start) * 1000
+            if resp.status == 200:
+                return ("OK", rtt_ms, None)
+            return ("FAIL", rtt_ms, f"HTTP {resp.status}")
+    except urllib.error.HTTPError as exc:
+        rtt_ms = (time.monotonic() - start) * 1000
+        return ("FAIL", rtt_ms, f"HTTP {exc.code}: {exc.reason}")
+    except (urllib.error.URLError, OSError) as exc:
+        rtt_ms = (time.monotonic() - start) * 1000
+        reason = getattr(exc, "reason", str(exc))
+        return ("FAIL", rtt_ms, f"Cannot connect to {url}: {reason}")
+
+
+# -------------------------------------------------------------------
+# Command
+# -------------------------------------------------------------------
+
+
+class PingCommand(BaseCommand):
+    """Ping LMCache or vLLM server (liveness check)."""
+
+    def name(self) -> str:
+        return "ping"
+
+    def help(self) -> str:
+        return "Ping LMCache or vLLM server (liveness check)."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.add_argument(
+            "target",
+            choices=["kvcache", "engine"],
+            help="What to ping.",
+        )
+        parser.add_argument(
+            "--url",
+            default=None,
+            help=(
+                "Server URL (default: http://localhost:8080 for kvcache, "
+                "http://localhost:8000 for engine)."
+            ),
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        base_url = normalize_url(args.url or DEFAULT_URLS[args.target])
+        endpoint = HEALTH_ENDPOINTS[args.target]
+        title = TITLES[args.target]
+
+        status, rtt_ms, error = ping(f"{base_url}{endpoint}")
+
+        metrics = self.create_metrics(title, args, width=30)
+        metrics.add("status", "Status", status)
+        metrics.add("round_trip_time_ms", "Round trip time (ms)", round(rtt_ms, 2))
+        metrics.emit()
+
+        if error:
+            print(error, file=sys.stderr)
+            sys.exit(1)

lmcache/cli/commands/query/__init__.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run one OpenAI-compatible inference request and report token/latency metrics."""
+
+# Standard
+import argparse
+import sys
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand, _add_output_args
+from lmcache.cli.commands.query.prompt import PromptBuilder
+from lmcache.cli.commands.query.request import Request
+
+
+class QueryCommand(BaseCommand):
+    """CLI command that sends one request to a serving engine."""
+
+    def name(self) -> str:
+        return "query"
+
+    def help(self) -> str:
+        return "Run one inference request and report metrics."
+
+    def add_arguments(self, _parser: argparse.ArgumentParser) -> None:
+        pass
+
+    def register(self, subparsers: argparse._SubParsersAction) -> None:
+        parser = subparsers.add_parser(
+            self.name(),
+            help=self.help(),
+            description=(
+                "Run one OpenAI-compatible inference request and report metrics."
+            ),
+        )
+        inner = parser.add_subparsers(
+            dest="query_target",
+            required=True,
+            metavar="{engine,kvcache}",
+        )
+        self._register_engine(inner)
+        self._register_kvcache(inner)
+
+    def _register_engine(self, subparsers: argparse._SubParsersAction) -> None:
+        parser = subparsers.add_parser(
+            "engine",
+            help="Send one request to an OpenAI-compatible HTTP API.",
+        )
+        parser.add_argument("--url", required=True, help="Serving engine base URL.")
+        parser.add_argument(
+            "--prompt",
+            required=True,
+            help="Prompt text with optional {name} placeholders.",
+        )
+        parser.add_argument(
+            "--model",
+            default=None,
+            metavar="ID",
+            help="Model ID for the serving engine.",
+        )
+        parser.add_argument(
+            "--max-tokens",
+            type=int,
+            default=128,
+            help="Maximum completion tokens (default: 128).",
+        )
+        parser.add_argument(
+            "--timeout",
+            type=float,
+            default=30.0,
+            help="HTTP timeout in seconds (default: 30).",
+        )
+        parser.add_argument(
+            "--documents",
+            action="extend",
+            nargs="+",
+            default=[],
+            metavar="NAME=PATH",
+            help=(
+                "Load file text for {NAME} in --prompt. "
+                "Accepts one or more NAME=PATH values."
+            ),
+        )
+        parser.add_argument(
+            "--path",
+            dest="documents",
+            action="extend",
+            nargs="+",
+            metavar="NAME=PATH",
+            help=argparse.SUPPRESS,
+        )
+        parser.add_argument(
+            "--completions",
+            action="store_true",
+            help="Use POST /v1/completions only.",
+        )
+        parser.add_argument(
+            "--chat-first",
+            action="store_true",
+            help="Try /v1/chat/completions first, then fall back to /v1/completions.",
+        )
+        _add_output_args(parser)
+        parser.set_defaults(func=self.execute)
+
+    def _register_kvcache(self, subparsers: argparse._SubParsersAction) -> None:
+        parser = subparsers.add_parser(
+            "kvcache",
+            help="Query KV-cache endpoints (not implemented yet).",
+        )
+        _add_output_args(parser)
+        parser.set_defaults(func=self.execute)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        handlers = {
+            "engine": self.query_engine,
+            "kvcache": self.query_kvcache,
+        }
+        handler = handlers.get(args.query_target)
+        if handler is None:
+            print(f"Unknown query target: {args.query_target}", file=sys.stderr)
+            sys.exit(1)
+        handler(args)
+
+    def query_engine(self, args: argparse.Namespace) -> None:
+        try:
+            prompt_builder = PromptBuilder(args.prompt, args.documents)
+            sender = Request(
+                base=args.url,
+                model=args.model,
+                max_tokens=args.max_tokens,
+                timeout=args.timeout,
+                completions_only=args.completions,
+                chat_first=args.chat_first,
+            )
+            engine_stats = sender.send_request(prompt_builder.complete_prompt)
+
+            model_id = args.model or str(engine_stats["model"][1])
+            metrics = self.create_metrics("Query Engine", args)
+            metrics.add("model", "Model", model_id)
+            prompt_name, prompt_value = engine_stats["prompt_tokens"]
+            metrics.add("prompt_tokens", prompt_name, int(prompt_value))
+            output_name, output_value = engine_stats["output_tokens"]
+            metrics.add("output_tokens", output_name, int(output_value))
+
+            latency = metrics.add_section("latency", "Latency Metrics")
+            for key, (name, value) in engine_stats.items():
+                if key in ("model", "prompt_tokens", "output_tokens"):
+                    continue
+                latency.add(key, name, round(float(value), 2))
+
+            metrics.emit()
+        except (RuntimeError, ValueError) as err:
+            print(str(err), file=sys.stderr)
+            sys.exit(1)
+
+    def query_kvcache(self, args: argparse.Namespace) -> None:
+        pass

lmcache/cli/commands/query/prompt.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Prompt placeholder expansion helpers for ``lmcache query``."""
+
+# Standard
+from pathlib import Path
+from typing import Optional
+import re
+
+PLACEHOLDER = re.compile(r"\{(\w+)\}")
+BUILTIN_DOCUMENT_PATHS = {
+    "lmcache": (
+        Path(__file__).resolve().parent.parent.parent / "documents" / "lmcache.txt"
+    ),
+}
+
+
+class PromptBuilder:
+    """Build a complete prompt from CLI inputs."""
+
+    def __init__(self, prompt: str, documents_args: Optional[list[str]] = None) -> None:
+        self._complete_prompt = self.expand_prompt(prompt, documents_args or [])
+
+    @property
+    def complete_prompt(self) -> str:
+        """Return the expanded prompt from constructor inputs."""
+        return self._complete_prompt
+
+    def expand_prompt(
+        self,
+        prompt: str,
+        documents_args: Optional[list[str]] = None,
+    ) -> str:
+        """Expand placeholders and return the complete prompt text.
+
+        Args:
+            prompt: The prompt template string with {name} placeholders.
+            documents_args: Optional list of document paths or NAME=PATH strings.
+
+        Returns:
+            The expanded prompt string.
+        """
+        docs_args = documents_args or []
+        documents, appended_text = resolve_documents(prompt, docs_args)
+
+        for key in _unique_placeholders(prompt):
+            if key not in documents:
+                unknown_documents(key)
+
+        complete_prompt = PLACEHOLDER.sub(
+            lambda match: documents[match.group(1)], prompt
+        )
+        if appended_text:
+            complete_prompt = (
+                f"{complete_prompt}\n{appended_text}"
+                if complete_prompt
+                else appended_text
+            )
+        return complete_prompt
+
+
+def resolve_documents(
+    prompt_template: str, documents_args: list[str]
+) -> tuple[dict[str, str], str]:
+    """Resolve ``--documents`` args into placeholder mapping and trailing text.
+
+    Args:
+        prompt_template: The prompt string containing placeholders.
+        documents_args: List of document arguments from the CLI.
+
+    Returns:
+        A tuple containing a dictionary mapping placeholder names to document
+        content and a string containing any appended document text.
+    """
+    placeholders = _unique_placeholders(prompt_template)
+    documents: dict[str, str] = {}
+    plain_docs: list[str] = []
+    for item in documents_args:
+        if "=" not in item:
+            plain_docs.append(_read_document_file(item))
+            continue
+        name, path = [x.strip() for x in item.split("=", 1)]
+        if not name:
+            raise ValueError(f"Invalid --documents {item!r}; empty name")
+        documents[name] = _read_document_file(path, name=name)
+
+    unresolved = [key for key in placeholders if key not in documents]
+    for idx, key in enumerate(unresolved[: len(plain_docs)]):
+        documents[key] = plain_docs[idx]
+    appended_docs = plain_docs[len(unresolved) :]
+
+    for key in placeholders:
+        if key in documents:
+            continue
+        builtin_path = BUILTIN_DOCUMENT_PATHS.get(key)
+        if builtin_path is None:
+            continue
+        documents[key] = _read_document_file(str(builtin_path), name=key)
+
+    return documents, "\n".join(appended_docs).strip()
+
+
+def unknown_documents(key: str) -> None:
+    """Raise an error for a missing documents placeholder.
+
+    Args:
+        key: The name of the missing placeholder.
+
+    Raises:
+        ValueError: Always raised with a descriptive message.
+    """
+    raise ValueError(
+        f"Unknown documents {key!r}. Define it with --documents {key}=PATH "
+    )
+
+
+def _read_document_file(path: str, *, name: Optional[str] = None) -> str:
+    file_path = Path(path).expanduser()
+    if not file_path.is_file():
+        if name is not None:
+            raise ValueError(f"documents file not found for {name!r}: {file_path}")
+        raise ValueError(f"documents file not found: {file_path}")
+    return file_path.read_text(encoding="utf-8", errors="replace")
+
+
+def _unique_placeholders(prompt_template: str) -> list[str]:
+    seen: set[str] = set()
+    keys: list[str] = []
+    for match in PLACEHOLDER.finditer(prompt_template):
+        key = match.group(1)
+        if key in seen:
+            continue
+        seen.add(key)
+        keys.append(key)
+    return keys