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

lmcache/cli/metrics/metrics.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Hierarchical metrics collector with handler-based output.
+
+Example usage::
+
+    from lmcache.cli.metrics import Metrics
+
+    metrics = Metrics(title="Bench KV Cache Result (30s)")
+
+    # Sectioned metrics
+    metrics.add_section("ops", "Operations (ops/s)")
+    metrics["ops"].add("store", "Store", 41.3)
+    metrics["ops"].add("retrieve", "Retrieve", 127.3)
+
+    # Top-level (flat) metrics
+    metrics.add("status", "Status", "OK")
+
+    metrics.emit()
+"""
+
+# Standard
+from typing import Any, Optional
+
+# First Party
+from lmcache.cli.metrics.handler import FileHandler, MetricsHandler
+from lmcache.cli.metrics.section import Section, sections_to_dict
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class Metrics:
+    """Hierarchical metrics collector with handler-based output.
+
+    Handlers are registered via :meth:`add_handler` and triggered
+    together by :meth:`emit`.  ``BaseCommand`` sets up default
+    handlers automatically, so command authors typically only need
+    to build metrics and call :meth:`emit`.
+
+    Args:
+        title: Report title shown in the header.
+    """
+
+    def __init__(self, title: str) -> None:
+        self._title = title
+        self._sections: list[Section] = []
+        self._section_map: dict[Optional[str], Section] = {}
+        self._handlers: list[MetricsHandler] = []
+
+    def title(self, title: str) -> None:
+        """Set the report title.
+
+        Args:
+            title: New report title shown in the header.
+        """
+        self._title = title
+
+    # -- Handler management -------------------------------------------------
+
+    def add_handler(self, handler: MetricsHandler) -> None:
+        """Register a handler.
+
+        Args:
+            handler: The handler to add.
+        """
+        self._handlers.append(handler)
+
+    # -- Section management -------------------------------------------------
+
+    def add_section(self, key: str, label: str) -> Section:
+        """Add a named section.
+
+        Args:
+            key: Machine-readable section key (used in JSON output and
+                for ``metrics["key"]`` access).
+            label: Human-readable label (used in terminal output).
+
+        Returns:
+            The newly created ``Section``.
+
+        Raises:
+            ValueError: If a section with the same *key* already exists.
+        """
+        if key in self._section_map:
+            raise ValueError(f"Section {key!r} already exists")
+        section = Section(key, label)
+        self._sections.append(section)
+        self._section_map[key] = section
+        return section
+
+    def add_list_section(self, group: str, key: str, label: str) -> Section:
+        """Add a section that belongs to a list group.
+
+        In terminal output, renders as a normal section with *label* as
+        the header. In JSON output, sections sharing the same *group*
+        are collected into a list: ``"group": [{...}, {...}]``.
+
+        Args:
+            group: The JSON key for the list (e.g., ``"models"``).
+            key: Unique machine key for this section instance.
+            label: Human-readable label (used in terminal output).
+
+        Returns:
+            The newly created ``Section``.
+        """
+        if key in self._section_map:
+            raise ValueError(f"Section {key!r} already exists")
+        section = Section(key, label, list_group=group)
+        self._sections.append(section)
+        self._section_map[key] = section
+        return section
+
+    def __getitem__(self, key: str) -> Section:
+        """Return the section registered under *key*.
+
+        Raises:
+            KeyError: If ``add_section(key, ...)`` was not called first.
+        """
+        return self._section_map[key]
+
+    # -- Flat (top-level) metrics -------------------------------------------
+
+    def _default_section(self) -> Section:
+        """Return the unnamed default section, creating it on first use."""
+        if None not in self._section_map:
+            section = Section(None, None)
+            # Insert at the beginning so flat metrics appear first
+            self._sections.insert(0, section)
+            self._section_map[None] = section
+        return self._section_map[None]
+
+    def add(self, key: str, label: str, value: Any) -> None:
+        """Record a top-level metric (not inside any named section).
+
+        Args:
+            key: Machine-readable key (used in JSON output).
+            label: Human-readable label (used in terminal output).
+            value: Metric value.
+        """
+        self._default_section().add(key, label, value)
+
+    # -- Output -------------------------------------------------------------
+
+    def emit(self) -> None:
+        """Trigger all registered handlers."""
+        for handler in self._handlers:
+            handler.emit(self._title, self._sections)
+        for handler in self._handlers:
+            if isinstance(handler, FileHandler):
+                logger.info("Results saved to %s", handler.path)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return metrics as a JSON-serialisable dictionary.
+
+        Returns:
+            A dict with ``"title"`` and ``"metrics"`` keys. Named
+            sections become nested dicts keyed by machine key. The
+            unnamed default section's entries are placed at the top
+            level of ``"metrics"``.
+        """
+        return sections_to_dict(self._title, self._sections)

lmcache/cli/metrics/section.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Section — a named group of metric entries."""
+
+# Standard
+from typing import Any, Optional
+
+
+class Section:
+    """A named group of metrics.
+
+    Each entry has a machine ``key`` (used in JSON), a human-readable
+    ``label`` (used in terminal output), and a ``value``.
+
+    Sections with the same :attr:`list_group` are collected into a
+    JSON list under that key (e.g., ``"models": [{...}, {...}]``).
+    In terminal output they render as normal independent sections.
+    """
+
+    def __init__(
+        self,
+        key: Optional[str],
+        label: Optional[str],
+        list_group: Optional[str] = None,
+    ) -> None:
+        self.key = key
+        self.label = label
+        self.list_group = list_group
+        self.entries: list[tuple[str, str, Any]] = []
+
+    def add(self, key: str, label: str, value: Any) -> None:
+        """Record a metric in this section.
+
+        Args:
+            key: Machine-readable key (used in JSON output).
+            label: Human-readable label (used in terminal output).
+            value: Metric value. Floats are formatted to 2 decimal
+                places on terminal output; strings are printed as-is.
+        """
+        self.entries.append((key, label, value))
+
+
+def sections_to_dict(
+    title: str,
+    sections: list[Section],
+) -> dict[str, Any]:
+    """Convert a title and sections to a JSON-serialisable dictionary.
+
+    Named sections become nested dicts keyed by machine key. The
+    unnamed default section's entries are placed at the top level
+    of ``"metrics"``.
+
+    Args:
+        title: The report title.
+        sections: Ordered list of ``Section`` objects.
+
+    Returns:
+        A dict with ``"title"`` and ``"metrics"`` keys.
+    """
+    metrics: dict[str, Any] = {}
+    list_groups: dict[str, list[dict[str, Any]]] = {}
+    for section in sections:
+        if section.key is None:
+            for key, _label, value in section.entries:
+                metrics[key] = value
+        elif section.list_group is not None:
+            section_dict: dict[str, Any] = {}
+            for key, _label, value in section.entries:
+                section_dict[key] = value
+            list_groups.setdefault(section.list_group, []).append(section_dict)
+        else:
+            section_dict = {}
+            for key, _label, value in section.entries:
+                section_dict[key] = value
+            metrics[section.key] = section_dict
+    for group_key, items in list_groups.items():
+        metrics[group_key] = items
+    return {"title": title, "metrics": metrics}

lmcache/connections.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# This file is copied from the vLLM project (https://github.com/vllm-project/vllm).
+# Original source file: [vllm/vllm/connections.py]
+# License: [Apache License 2.0]
+# Modifications: header name
+
+# Standard
+from pathlib import Path
+from typing import Mapping, MutableMapping, Optional
+from urllib.parse import urlparse
+
+# Third Party
+import aiohttp
+import requests
+
+
+class HTTPConnection:
+    """Helper class to send HTTP requests."""
+
+    def __init__(self, *, reuse_client: bool = True) -> None:
+        super().__init__()
+
+        self.reuse_client = reuse_client
+
+        self._sync_client: Optional[requests.Session] = None
+        self._async_client: Optional[aiohttp.ClientSession] = None
+
+    def get_sync_client(self) -> requests.Session:
+        if self._sync_client is None or not self.reuse_client:
+            self._sync_client = requests.Session()
+
+        return self._sync_client
+
+    # NOTE: We intentionally use an async function even though it is not
+    # required, so that the client is only accessible inside async event loop
+    async def get_async_client(self) -> aiohttp.ClientSession:
+        if self._async_client is None or not self.reuse_client:
+            self._async_client = aiohttp.ClientSession()
+
+        return self._async_client
+
+    def _validate_http_url(self, url: str):
+        parsed_url = urlparse(url)
+
+        if parsed_url.scheme not in ("http", "https"):
+            raise ValueError(
+                "Invalid HTTP URL: A valid HTTP URL must have scheme 'http' or 'https'."
+            )
+
+    def _headers(self, **extras: str) -> MutableMapping[str, str]:
+        return {"User-Agent": "LMCache", **extras}
+
+    def get_response(
+        self,
+        url: str,
+        *,
+        stream: bool = False,
+        timeout: Optional[float] = None,
+        extra_headers: Optional[Mapping[str, str]] = None,
+    ):
+        self._validate_http_url(url)
+
+        client = self.get_sync_client()
+        extra_headers = extra_headers or {}
+
+        return client.get(
+            url,
+            headers=self._headers(**extra_headers),
+            stream=stream,
+            timeout=timeout,
+        )
+
+    async def get_async_response(
+        self,
+        url: str,
+        *,
+        timeout: Optional[float] = None,
+        extra_headers: Optional[Mapping[str, str]] = None,
+    ):
+        self._validate_http_url(url)
+
+        client = await self.get_async_client()
+        extra_headers = extra_headers or {}
+
+        return client.get(url, headers=self._headers(**extra_headers), timeout=timeout)
+
+    def get_bytes(self, url: str, *, timeout: Optional[float] = None) -> bytes:
+        with self.get_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            return r.content
+
+    async def async_get_bytes(
+        self,
+        url: str,
+        *,
+        timeout: Optional[float] = None,
+    ) -> bytes:
+        async with await self.get_async_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            return await r.read()
+
+    def get_text(self, url: str, *, timeout: Optional[float] = None) -> str:
+        with self.get_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            return r.text
+
+    async def async_get_text(
+        self,
+        url: str,
+        *,
+        timeout: Optional[float] = None,
+    ) -> str:
+        async with await self.get_async_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            return await r.text()
+
+    def get_json(self, url: str, *, timeout: Optional[float] = None) -> str:
+        with self.get_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            return r.json()
+
+    async def async_get_json(
+        self,
+        url: str,
+        *,
+        timeout: Optional[float] = None,
+    ) -> str:
+        async with await self.get_async_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            return await r.json()
+
+    def download_file(
+        self,
+        url: str,
+        save_path: Path,
+        *,
+        timeout: Optional[float] = None,
+        chunk_size: int = 128,
+    ) -> Path:
+        with self.get_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            with save_path.open("wb") as f:
+                for chunk in r.iter_content(chunk_size):
+                    f.write(chunk)
+
+        return save_path
+
+    async def async_download_file(
+        self,
+        url: str,
+        save_path: Path,
+        *,
+        timeout: Optional[float] = None,
+        chunk_size: int = 128,
+    ) -> Path:
+        async with await self.get_async_response(url, timeout=timeout) as r:
+            r.raise_for_status()
+
+            with save_path.open("wb") as f:
+                async for chunk in r.content.iter_chunked(chunk_size):
+                    f.write(chunk)
+
+        return save_path
+
+
+global_http_connection = HTTPConnection()

lmcache/integration/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+

lmcache/integration/base_service_factory.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+BaseServiceFactory: Abstract interface for creating LMCache service components.
+
+Each serving engine integration (e.g., vLLM) should implement a concrete
+ServiceFactory that determines which components to create for each role.
+"""
+
+# Standard
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Optional, Union
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.health_monitor.base import HealthMonitor
+from lmcache.v1.health_monitor.constants import (
+    DEFAULT_PING_INTERVAL,
+    PING_INTERVAL_CONFIG_KEY,
+)
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.observability import PrometheusLogger
+    from lmcache.v1.cache_engine import LMCacheEngine
+    from lmcache.v1.internal_api_server.api_server import InternalAPIServer
+    from lmcache.v1.lookup_client.abstract_client import LookupClientInterface
+    from lmcache.v1.lookup_client.lmcache_async_lookup_client import (
+        LMCacheAsyncLookupServer,
+    )
+    from lmcache.v1.lookup_client.lmcache_lookup_client import LMCacheLookupServer
+    from lmcache.v1.manager import LMCacheManager
+    from lmcache.v1.metadata import LMCacheMetadata
+    from lmcache.v1.offload_server.zmq_server import ZMQOffloadServer
+    from lmcache.v1.plugin.runtime_plugin_launcher import RuntimePluginLauncher
+
+logger = init_logger(__name__)
+
+
+class BaseServiceFactory(ABC):
+    """Abstract base for creating LMCache service components.
+
+    Subclasses must implement all methods to provide the appropriate
+    components for their serving engine integration.
+    """
+
+    @abstractmethod
+    def get_engine_instance_id(self) -> str:
+        """Return the instance_id used to register the engine with
+        LMCacheEngineBuilder. Used by LMCacheManager for engine destruction."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_or_create_metadata(self) -> Optional["LMCacheMetadata"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_or_create_lmcache_engine(self) -> Optional["LMCacheEngine"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_lookup_client(self) -> Optional["LookupClientInterface"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_prometheus_logger(self) -> Optional["PrometheusLogger"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_lookup_server(
+        self,
+    ) -> Optional[Union["LMCacheLookupServer", "LMCacheAsyncLookupServer"]]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_offload_server(self) -> Optional["ZMQOffloadServer"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_runtime_plugin_launcher(
+        self,
+    ) -> Optional["RuntimePluginLauncher"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_internal_api_server(
+        self, lmcache_manager: "LMCacheManager"
+    ) -> Optional["InternalAPIServer"]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def maybe_create_health_monitor(
+        self, lmcache_manager: "LMCacheManager"
+    ) -> Optional[HealthMonitor]:
+        raise NotImplementedError
+
+    def _create_health_monitor(
+        self,
+        lmcache_manager: "LMCacheManager",
+        config: LMCacheEngineConfig,
+        engine: Optional["LMCacheEngine"] = None,
+    ) -> HealthMonitor:
+        """Create, configure, and start the health monitor.
+
+        Shared implementation used by subclass maybe_create_health_monitor.
+        """
+        # First Party
+        from lmcache.observability import PrometheusLogger
+        from lmcache.v1.periodic_thread import (
+            PeriodicThreadRegistry,
+            ThreadLevel,
+        )
+
+        ping_interval = config.get_extra_config_value(
+            PING_INTERVAL_CONFIG_KEY, DEFAULT_PING_INTERVAL
+        )
+        health_monitor = HealthMonitor(
+            manager=lmcache_manager,
+            ping_interval=ping_interval,
+        )
+
+        if engine is not None:
+            engine.set_health_monitor(health_monitor)
+
+        health_monitor.start()
+        logger.info("Health monitor initialized and started")
+
+        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
+        if prometheus_logger is not None:
+            prometheus_logger.lmcache_is_healthy.set_function(
+                lambda: 1 if lmcache_manager.is_healthy() else 0
+            )
+
+            registry = PeriodicThreadRegistry.get_instance()
+
+            prometheus_logger.periodic_threads_total_count.set_function(
+                lambda: len(registry.get_all())
+            )
+            prometheus_logger.periodic_threads_running_count.set_function(
+                lambda: registry.get_running_count()
+            )
+            prometheus_logger.periodic_threads_active_count.set_function(
+                lambda: registry.get_active_count()
+            )
+
+            for level in ThreadLevel:
+                level_name = level.value
+                total_attr = f"periodic_threads_{level_name}_total"
+                running_attr = f"periodic_threads_{level_name}_running"
+                active_attr = f"periodic_threads_{level_name}_active"
+
+                if hasattr(prometheus_logger, total_attr):
+                    getattr(prometheus_logger, total_attr).set_function(
+                        lambda lvl=level: registry.get_count_by_level(lvl)["total"]
+                    )
+                if hasattr(prometheus_logger, running_attr):
+                    getattr(prometheus_logger, running_attr).set_function(
+                        lambda lvl=level: registry.get_count_by_level(lvl)["running"]
+                    )
+                if hasattr(prometheus_logger, active_attr):
+                    getattr(prometheus_logger, active_attr).set_function(
+                        lambda lvl=level: registry.get_count_by_level(lvl)["active"]
+                    )
+
+        return health_monitor

lmcache/integration/request_telemetry/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

lmcache/integration/request_telemetry/base.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Abstract base class for request telemetry.
+
+This module provides the interface for tracking request-level events
+in LMCache, such as when a request finishes and its associated async
+save operations complete.
+"""
+
+# Standard
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class RequestTelemetry(ABC):
+    """
+    Abstract base class for request telemetry.
+
+    This class defines the interface for capturing request-level telemetry
+    events. Implementations can log events, emit metrics, or perform other
+    actions when specific request lifecycle events occur.
+    """
+
+    @abstractmethod
+    def __init__(self, config: dict[str, Any]) -> None:
+        pass
+
+    @abstractmethod
+    def on_request_store_finished(
+        self,
+        request_ids_set: set[str],
+        model_name: str,
+        world_size: int,
+        kv_rank: int,
+    ) -> None:
+        """
+        Callback when request finishes AND all its KV cache store ops completes.
+
+        This method ensures that request_ids_set is not empty.
+
+        Technically this function is implemented by inspecting the return value
+        of `get_finished` method.
+        """
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        pass
+
+    def __del__(self) -> None:
+        self.close()