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

lmcache/v1/health_monitor/checks/remote_backend_check.py

@@ -0,0 +1,304 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Health check for RemoteBackend.
+"""
+
+# Standard
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, List, Optional
+import asyncio
+import time
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.observability import LMCStatsMonitor
+from lmcache.utils import CacheEngineKey
+from lmcache.v1.health_monitor.base import HealthCheck
+from lmcache.v1.health_monitor.constants import (
+    DEFAULT_FALLBACK_POLICY,
+    DEFAULT_GET_BLOCKING_FAILED_THRESHOLD,
+    DEFAULT_PING_TIMEOUT,
+    DEFAULT_WAITING_TIME_FOR_RECOVERY,
+    FALLBACK_POLICY_CONFIG_KEY,
+    GET_BLOCKING_FAILED_THRESHOLD_CONFIG_KEY,
+    PING_GENERIC_ERROR_CODE,
+    PING_TIMEOUT_CONFIG_KEY,
+    PING_TIMEOUT_ERROR_CODE,
+    WAITING_TIME_FOR_RECOVERY_CONFIG_KEY,
+    FallbackPolicy,
+)
+from lmcache.v1.storage_backend.connector import InstrumentedRemoteConnector
+from lmcache.v1.storage_backend.connector.audit_connector import AuditConnector
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.manager import LMCacheManager
+    from lmcache.v1.storage_backend.remote_backend import RemoteBackend
+
+logger = init_logger(__name__)
+
+
+class RemoteBackendHealthCheck(HealthCheck):
+    """
+    Health check for RemoteBackend by pinging the remote connector.
+
+    This check verifies that the remote backend is reachable and responsive
+    by sending periodic ping requests.
+
+    Fallback Policies:
+        - RECOMPUTE (default): Skip all cache operations when
+          remote backend is unhealthy
+        - LOCAL_CPU: Bypass remote backend and use local CPU with hot_cache enabled
+    """
+
+    def __init__(
+        self,
+        backend: "RemoteBackend",
+    ):
+        self.backend = backend
+        # Get fallback policy from config
+        fallback_policy_str = backend.config.get_extra_config_value(
+            FALLBACK_POLICY_CONFIG_KEY, DEFAULT_FALLBACK_POLICY.value
+        )
+        # Convert string to FallbackPolicy enum
+        if isinstance(fallback_policy_str, str):
+            try:
+                self._fallback_policy = FallbackPolicy(fallback_policy_str)
+            except ValueError:
+                logger.warning(
+                    f"Invalid fallback_policy '{fallback_policy_str}' "
+                    f"for {backend}, using default: "
+                    f"{DEFAULT_FALLBACK_POLICY}"
+                )
+                self._fallback_policy = DEFAULT_FALLBACK_POLICY
+        elif isinstance(fallback_policy_str, FallbackPolicy):
+            self._fallback_policy = fallback_policy_str
+        self.failure_time: Optional[float] = None
+        self._stats_monitor = LMCStatsMonitor.GetOrCreate()
+        self._backend_name: Optional[str] = None
+        self._last_get_blocking_failed_count = 0
+
+    @classmethod
+    def create(cls, manager: "LMCacheManager") -> List[HealthCheck]:
+        """
+        Create RemoteBackendHealthCheck instances from a LMCacheManager.
+
+        This method finds all RemoteBackend instances in the storage manager
+        and creates a health check for each one.
+
+        Args:
+            manager: The LMCacheManager instance
+
+        Returns:
+            List of RemoteBackendHealthCheck instances
+        """
+        # Import here to avoid circular imports
+        # First Party
+        from lmcache.v1.storage_backend.remote_backend import RemoteBackend
+
+        instances: List[HealthCheck] = []
+
+        # Get engine from manager
+        engine = manager.lmcache_engine
+        if engine is None or engine.storage_manager is None:
+            return instances
+
+        for backend_name, backend in engine.storage_manager.storage_backends.items():
+            if isinstance(backend, RemoteBackend):
+                check = cls(backend)
+                check._backend_name = backend_name
+                instances.append(check)
+                logger.info(f"Created {check} for {backend_name}")
+
+        return instances
+
+    def name(self) -> str:
+        return f"RemoteBackendHealthCheck({self.backend.remote_url})"
+
+    @property
+    def fallback_policy(self) -> FallbackPolicy:
+        """Return the fallback policy for this health check."""
+        return self._fallback_policy
+
+    def get_bypass_backend_name(self) -> Optional[str]:
+        """
+        Return the backend name to bypass when this health check fails.
+
+        Returns:
+            Optional[str]: The backend name (e.g., "RemoteBackend")
+        """
+        return self._backend_name
+
+    def _get_ping_timeout(self) -> float:
+        """Get the ping timeout from the backend config."""
+        return self.backend.config.get_extra_config_value(
+            PING_TIMEOUT_CONFIG_KEY, DEFAULT_PING_TIMEOUT
+        )
+
+    def _try_reinitialize_connection(self) -> bool:
+        """
+        Try to reinitialize the connection if connector is None.
+
+        Returns:
+            bool: True if connection was successfully initialized, False otherwise
+        """
+        if self.backend.connection is not None:
+            return True
+
+        logger.warning("Connector is None, re-initializing connection.")
+        self.backend.init_connection()
+
+        return self.backend.connection is not None
+
+    def check(self) -> bool:
+        """
+        Perform a health check on remote backend, which includes the following checks:
+
+        - get_blocking/batched_get_blocking, if failed count >= threshold,
+        which means check failed, update failure_time and return False.
+        If failure_time is not None, wait for more than`waiting_time_for_recovery`
+        seconds before resuming the check.
+
+        - ping, if connector supports ping, send a ping request to remote connector.
+
+        Returns:
+            bool: True if all checks succeeds, False otherwise
+        """
+        # Try to reinitialize connection if needed
+        if not self._try_reinitialize_connection():
+            return False
+
+        # At this point, connector is guaranteed to be not None
+        connector = self.backend.connection
+        assert connector is not None
+
+        if self.failure_time is not None:
+            waiting_time = self.backend.config.get_extra_config_value(
+                WAITING_TIME_FOR_RECOVERY_CONFIG_KEY,
+                DEFAULT_WAITING_TIME_FOR_RECOVERY,
+            )
+            if (
+                time.time() - self.failure_time > waiting_time
+                and self._put_and_get_check()
+            ):
+                # recover from get blocking failed
+                logger.info(
+                    "Failure time: %s, current time: %s, "
+                    "recover from get blocking failed",
+                    self.failure_time,
+                    time.time(),
+                )
+                self.failure_time = None
+            else:
+                logger.info(
+                    "Failure time: %s, current time: %s, "
+                    "still in get blocking failed recovery window",
+                    self.failure_time,
+                    time.time(),
+                )
+                return False
+
+        # Check read failed
+        current_get_blocking_failed_count = self.backend.get_blocking_failed_count
+        get_blocking_failed_count = (
+            current_get_blocking_failed_count - self._last_get_blocking_failed_count
+        )
+        self._last_get_blocking_failed_count = current_get_blocking_failed_count
+        threshold = self.backend.config.get_extra_config_value(
+            GET_BLOCKING_FAILED_THRESHOLD_CONFIG_KEY,
+            DEFAULT_GET_BLOCKING_FAILED_THRESHOLD,
+        )
+        if get_blocking_failed_count >= threshold:
+            logger.warning(
+                "Detected %s get blocking failed in interval, threshold: %s",
+                get_blocking_failed_count,
+                threshold,
+            )
+            self.failure_time = time.time()
+            return False
+
+        # If connector doesn't support ping, assume it's healthy
+        if not connector.support_ping():
+            return True
+
+        # Check ping
+        try:
+            start_time = time.perf_counter()
+            future = asyncio.run_coroutine_threadsafe(
+                connector.ping(), self.backend.loop
+            )
+            error_code = future.result(timeout=self._get_ping_timeout())
+            latency = (time.perf_counter() - start_time) * 1000
+
+            # Record ping latency
+            self._stats_monitor.update_remote_ping_latency(latency)
+            # Record error code (0 means success)
+            self._stats_monitor.update_remote_ping_error_code(error_code)
+
+            if error_code != 0:
+                logger.warning(f"Ping failed with error code: {error_code}")
+                return False
+
+            return True
+
+        except asyncio.TimeoutError:
+            logger.warning("Ping timeout")
+            self._stats_monitor.update_remote_ping_error_code(PING_TIMEOUT_ERROR_CODE)
+            return False
+        except Exception as e:
+            logger.error(f"Ping error: {e}")
+            self._stats_monitor.update_remote_ping_error_code(PING_GENERIC_ERROR_CODE)
+            return False
+
+    def _put_and_get_check(self) -> bool:
+        if self.backend.local_cpu_backend is None or self.backend.connection is None:
+            return False
+
+        with self._resource_manager() as (put_obj, get_obj):
+            if put_obj is None:
+                return False
+            if get_obj is None:
+                logger.warning("Get failed, the return value is None, check failed.")
+                return False
+            return torch.equal(put_obj.raw_tensor, get_obj.raw_tensor)
+
+    @contextmanager
+    def _resource_manager(self):
+        key = CacheEngineKey(
+            model_name="test",
+            world_size=1,
+            worker_id=0,
+            chunk_hash=0,
+            dtype=torch.bfloat16,
+        )
+        connector = self.backend.connection
+        if isinstance(connector, InstrumentedRemoteConnector):
+            connector = connector.getWrappedConnector()
+            if isinstance(connector, AuditConnector):
+                connector = connector.real_connector
+        shapes = connector.meta_shapes
+        dtypes = connector.meta_dtypes
+        fmt = connector.meta_fmt
+        put_obj, get_obj = None, None
+        try:
+            # put
+            put_obj = self.backend.local_cpu_backend.allocate(shapes, dtypes, fmt)
+            future = self.backend.submit_put_task(key, put_obj)
+            future.result(timeout=self._get_ping_timeout())
+            # get
+            get_obj = self.backend.get_blocking(key)
+            yield put_obj, get_obj
+        except asyncio.TimeoutError:
+            logger.warning("Put timeout, check failed.")
+            yield None, None
+        except Exception as e:
+            logger.error(f"Put error, check failed: {e}")
+            yield None, None
+        finally:
+            if put_obj is not None:
+                put_obj.ref_count_down()
+            if get_obj is not None:
+                get_obj.ref_count_down()

lmcache/v1/health_monitor/constants.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Constants for health monitoring.
+"""
+
+# Standard
+from enum import Enum
+
+
+class FallbackPolicy(str, Enum):
+    """Fallback policy when health check fails."""
+
+    RECOMPUTE = "recompute"  # Skip all cache operations, fall back to recomputation
+    LOCAL_CPU = "local_cpu"  # Fall back to local CPU backend only
+
+
+# Ping error codes
+PING_TIMEOUT_ERROR_CODE = -1
+PING_GENERIC_ERROR_CODE = -2
+
+# Configuration keys
+PING_TIMEOUT_CONFIG_KEY = "ping_timeout"
+PING_INTERVAL_CONFIG_KEY = "ping_interval"
+FALLBACK_POLICY_CONFIG_KEY = "fallback_policy"
+GET_BLOCKING_FAILED_THRESHOLD_CONFIG_KEY = "get_blocking_failed_threshold"
+WAITING_TIME_FOR_RECOVERY_CONFIG_KEY = "waiting_time_for_recovery"
+
+# Default values
+DEFAULT_PING_TIMEOUT = 5.0
+DEFAULT_PING_INTERVAL = 30.0
+DEFAULT_FALLBACK_POLICY = FallbackPolicy.RECOMPUTE
+DEFAULT_GET_BLOCKING_FAILED_THRESHOLD = 10
+DEFAULT_WAITING_TIME_FOR_RECOVERY = 300.0
+
+# Memory thresholds
+DEFAULT_MEMORY_THRESHOLD_PERCENT = 95.0  # Unhealthy if memory usage > 95%

lmcache/v1/internal_api_server/api_registry.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from pathlib import Path
+from typing import List, Literal, Optional
+import importlib.util
+import pkgutil
+
+# Third Party
+from fastapi import APIRouter
+
+APICategory = Literal["common", "vllm", "controller"]
+
+
+class APIRegistry:
+    """
+    Automatically discovers and registers API routes by category
+
+    Categories:
+    - common: APIs that work for all components (metrics, logs, etc.)
+    - vllm: APIs specific to vLLM scheduler/worker
+    - controller: APIs specific to LMCache controller
+    """
+
+    def __init__(self, app):
+        self.app = app
+        self.router = APIRouter()
+
+    def register_all_apis(self, categories: Optional[List[APICategory]] = None):
+        """
+        Discover and register API modules from specified categories
+
+        Args:
+            categories: List of categories to register.
+                       If None, registers all categories.
+        """
+        if categories is None:
+            categories = ["common", "vllm", "controller"]
+
+        package_path = Path(__file__).parent
+        package_name = __package__
+
+        for category in categories:
+            category_path = package_path / category
+            if not category_path.exists():
+                continue
+
+            category_package = f"{package_name}.{category}"
+
+            for _, module_name, _ in pkgutil.iter_modules([str(category_path)]):
+                if module_name.endswith("_api"):
+                    full_module_name = f"{category_package}.{module_name}"
+                    module = importlib.import_module(full_module_name)
+                    # Include the router if it exists
+                    if hasattr(module, "router") and isinstance(
+                        module.router, APIRouter
+                    ):
+                        self.router.include_router(module.router)
+
+        self.app.include_router(self.router)

lmcache/v1/internal_api_server/api_server.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING
+import asyncio
+import os
+import threading
+
+# Third Party
+from fastapi import FastAPI
+import uvicorn
+
+# First Party
+from lmcache.logging import init_logger
+
+# Local
+from .api_registry import APIRegistry
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.manager import LMCacheManager
+
+logger = init_logger(__name__)
+
+app = FastAPI()
+
+# Automatically register common, vllm, and controller APIs
+registry = APIRegistry(app)
+registry.register_all_apis(categories=["common", "vllm", "controller"])
+
+
+class InternalAPIServer:
+    def __init__(self, lmcache_manager: "LMCacheManager"):
+        lmcache_engine = lmcache_manager.lmcache_engine
+
+        # Check if lmcache_engine is None and handle accordingly
+        if lmcache_engine is None:
+            # Use manager's config directly when engine is not available
+            config = lmcache_manager.config
+            port_offset = 0  # Default for scheduler mode
+        else:
+            config = lmcache_engine.config
+            # 0 for scheduler, 1 for worker 0, 2 for worker 1, ...
+            port_offset = 1 + lmcache_engine.metadata.worker_id
+
+        self.port = config.internal_api_server_port_start + port_offset
+        self.socket_path_prefix = config.internal_api_server_socket_path_prefix
+        self.socket_path = (
+            f"{self.socket_path_prefix}_{self.port}"
+            if self.socket_path_prefix
+            else None
+        )
+        include_index_list = config.internal_api_server_include_index_list
+
+        self.enable = True
+        if not config.internal_api_server_enabled or (
+            include_index_list and port_offset not in include_index_list
+        ):
+            logger.info(
+                f"Internal API server disabled. internal_api_server_enabled="
+                f"{config.internal_api_server_enabled}, port_offset={port_offset}, "
+                f"port={self.port}, socket_path={self.socket_path}, "
+                f"include_index_list={include_index_list}"
+            )
+            self.enable = False
+            return
+
+        uvicorn_config = {
+            "app": app,
+            "host": config.internal_api_server_host,
+            "loop": "uvloop",
+            "http": "httptools",
+            "access_log": config.get_extra_config_value(
+                "internal_api_server_access_log", True
+            ),
+            "log_level": config.get_extra_config_value(
+                "internal_api_server_log_level", "warning"
+            ),
+        }
+
+        if self.socket_path:
+            self.server_log_info = f"socket {self.socket_path}"
+            logger.info(f"Init internal API server on {self.server_log_info}")
+            uvicorn_config["uds"] = self.socket_path
+            # Ensure socket directory exists
+            os.makedirs(os.path.dirname(self.socket_path), exist_ok=True)
+            # Remove existing socket file if exists
+            if os.path.exists(self.socket_path):
+                os.unlink(self.socket_path)
+        else:
+            self.server_log_info = f"port {self.port}"
+            logger.info(f"Init internal API server on {self.server_log_info}")
+            uvicorn_config["port"] = self.port
+
+        self.server = uvicorn.Server(uvicorn.Config(**uvicorn_config))
+        app.state.lmcache_adapter = lmcache_manager
+
+    async def run(self):
+        logger.info(f"Running LMCache internal API server on {self.server_log_info}")
+        if self.server:
+            await self.server.serve()
+
+    def start(self):
+        if not self.enable:
+            return
+        logger.info(f"Starting LMCache internal API server on {self.server_log_info}")
+        threading.Thread(
+            target=asyncio.run,
+            args=(self.run(),),
+            daemon=True,
+            name="api-server-thread",
+        ).start()
+
+    def stop(self):
+        if not self.enable:
+            return
+        logger.info("Stopping LMCache internal API server")
+        if self.server:
+            self.server.should_exit = True
+            if self.socket_path and os.path.exists(self.socket_path):
+                os.unlink(self.socket_path)

lmcache/v1/internal_api_server/common/__init__.py

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

lmcache/v1/internal_api_server/common/env_api.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import json
+import os
+
+# Third Party
+from fastapi import APIRouter
+from starlette.responses import PlainTextResponse
+
+router = APIRouter()
+
+
+@router.get("/env")
+async def get_env():
+    """
+    Get all environment variables
+    """
+    env_dict = dict(os.environ)
+    return PlainTextResponse(
+        content=json.dumps(env_dict, indent=2, sort_keys=True),
+        media_type="text/plain",
+    )

lmcache/v1/internal_api_server/common/loglevel_api.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional
+import logging
+
+# Third Party
+from fastapi import APIRouter
+from starlette.responses import PlainTextResponse
+
+router = APIRouter()
+
+
+@router.get("/loglevel")
+async def get_or_set_log_level(
+    logger_name: Optional[str] = None, level: Optional[str] = None
+):
+    """
+    Get or set the log level for a logger.
+    - No parameters: List all loggers and their levels.
+    - With logger_name: Get the level of the specified logger.
+    - With logger_name and level: Set the level of the specified logger.
+    """
+    if not logger_name and not level:
+        # List all loggers and their levels
+        loggers = logging.Logger.manager.loggerDict
+        result = "=== Loggers and Levels ===\n"
+        for name, logger_obj in loggers.items():
+            if isinstance(logger_obj, logging.Logger):
+                result += f"{name}: {logging.getLevelName(logger_obj.level)}\n"
+        return PlainTextResponse(content=result, media_type="text/plain")
+    elif logger_name and not level:
+        # Get the level of the specified logger
+        target_logger = logging.getLogger(logger_name)
+        return PlainTextResponse(
+            content=f"{logger_name}: {logging.getLevelName(target_logger.level)}",
+            media_type="text/plain",
+        )
+    elif logger_name and level:
+        # Set the level of the specified logger
+        target_logger = logging.getLogger(logger_name)
+        try:
+            level_value = getattr(logging, level.upper())
+            target_logger.setLevel(level_value)
+            # Set the level of all handlers
+            for handler in target_logger.handlers:
+                handler.setLevel(level_value)
+            return PlainTextResponse(
+                content=f"Set {logger_name} level to {level.upper()} "
+                "(including all handlers)",
+                media_type="text/plain",
+            )
+        except AttributeError:
+            return PlainTextResponse(
+                content=f"Invalid log level: {level}",
+                media_type="text/plain",
+                status_code=400,
+            )

lmcache/v1/internal_api_server/common/metrics_api.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# Third Party
+from fastapi import APIRouter
+from prometheus_client import REGISTRY, generate_latest
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+
+# First Party
+from lmcache.observability import reset_observability_metrics
+
+router = APIRouter()
+
+
+@router.get("/metrics")
+async def get_metrics(request: Request):
+    """
+    Provide Prometheus metrics data
+    """
+    metrics_data = generate_latest(REGISTRY)
+    return PlainTextResponse(content=metrics_data, media_type="text/plain")
+
+
+@router.post("/metrics/reset")
+async def reset_metrics():
+    """
+    Reset Prometheus metrics to their initial state.
+    """
+    reset_observability_metrics()
+    return PlainTextResponse(content="ok", media_type="text/plain")