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

lmcache/v1/internal_api_server/vllm/chunk_statistics_api.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Tuple, Union
+import json
+import time
+
+# Third Party
+from fastapi import APIRouter
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+
+# First Party
+from lmcache.v1.lookup_client.abstract_client import LookupClientInterface
+from lmcache.v1.lookup_client.chunk_statistics_lookup_client import (
+    ChunkStatisticsLookupClient,
+)
+
+router = APIRouter()
+
+
+def _create_json_response(data: dict, status_code: int = 200) -> PlainTextResponse:
+    return PlainTextResponse(
+        content=json.dumps(data, indent=2),
+        media_type="application/json",
+        status_code=status_code,
+    )
+
+
+def _get_lookup_client(
+    request: Request,
+) -> Tuple[
+    Optional[Union[LookupClientInterface, ChunkStatisticsLookupClient]],
+    Optional[PlainTextResponse],
+]:
+    lookup_client = getattr(request.app.state.lmcache_adapter, "lookup_client", None)
+    if not lookup_client:
+        return None, _create_json_response(
+            {"error": "API unavailable", "message": "Lookup client not configured."},
+            status_code=503,
+        )
+    return lookup_client, None
+
+
+def _get_statistics_client(
+    request: Request,
+) -> Tuple[Optional[ChunkStatisticsLookupClient], Optional[PlainTextResponse]]:
+    lookup_client, error_response = _get_lookup_client(request)
+    if error_response:
+        return None, error_response
+    if not isinstance(lookup_client, ChunkStatisticsLookupClient):
+        return None, _create_json_response(
+            {
+                "error": "Not available",
+                "message": "Client does not support statistics.",
+            },
+            status_code=400,
+        )
+    return lookup_client, None
+
+
+def _handle_exception(operation: str, error: Exception) -> PlainTextResponse:
+    return _create_json_response(
+        {"error": f"Failed to {operation}", "message": str(error)}, status_code=500
+    )
+
+
+@router.post("/chunk_statistics/start")
+async def start_chunk_statistics(request: Request):
+    try:
+        lookup_client, error_response = _get_lookup_client(request)
+        if error_response:
+            return error_response
+        assert lookup_client is not None
+        if isinstance(lookup_client, ChunkStatisticsLookupClient):
+            lookup_client.start_statistics()
+            return _create_json_response({"status": "success", "message": "Started"})
+        return _create_json_response(
+            {
+                "error": "Not available",
+                "message": "Client does not support statistics.",
+            },
+            status_code=400,
+        )
+    except Exception as e:
+        return _handle_exception("start", e)
+
+
+@router.post("/chunk_statistics/stop")
+async def stop_chunk_statistics(request: Request):
+    try:
+        stats_client, error_response = _get_statistics_client(request)
+        if error_response:
+            return error_response
+        assert stats_client is not None
+        stats_client.stop_statistics()
+        return _create_json_response({"status": "success", "message": "Stopped"})
+    except Exception as e:
+        return _handle_exception("stop", e)
+
+
+@router.post("/chunk_statistics/reset")
+async def reset_chunk_statistics(request: Request):
+    try:
+        stats_client, error_response = _get_statistics_client(request)
+        if error_response:
+            return error_response
+        assert stats_client is not None
+        stats_client.reset_statistics()
+        return _create_json_response({"status": "success", "message": "Reset"})
+    except Exception as e:
+        return _handle_exception("reset", e)
+
+
+@router.get("/chunk_statistics/status")
+async def get_chunk_statistics_status(request: Request):
+    try:
+        stats_client, error_response = _get_statistics_client(request)
+        if error_response:
+            return error_response
+        assert stats_client is not None
+        config = request.app.state.lmcache_adapter.config
+        stats = stats_client.get_statistics()
+        stats.update(
+            {
+                "timestamp": time.time(),
+                "auto_exit_enabled": (
+                    config.chunk_statistics_auto_exit_timeout_hours > 0.0
+                    or config.chunk_statistics_auto_exit_target_unique_chunks
+                    is not None
+                ),
+                "auto_exit_timeout_hours": (
+                    config.chunk_statistics_auto_exit_timeout_hours
+                ),
+                "auto_exit_target_unique_chunks": (
+                    config.chunk_statistics_auto_exit_target_unique_chunks
+                ),
+            }
+        )
+        return _create_json_response(stats)
+    except Exception as e:
+        return _handle_exception("get status", e)

lmcache/v1/internal_api_server/vllm/conf_api.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Any, Dict, Optional, Union
+import json
+
+# Third Party
+from fastapi import APIRouter
+from starlette.requests import Request
+from starlette.responses import JSONResponse, PlainTextResponse
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.config import _CONFIG_DEFINITIONS, LMCacheEngineConfig
+from lmcache.v1.config_base import validate_and_set_config_value
+
+logger = init_logger(__name__)
+
+router = APIRouter()
+
+
+def _is_mutable_config(key: str) -> bool:
+    """Check if a config key is mutable at runtime.
+
+    NOTE: This is currently experimental. All configs default to
+    mutable=True unless explicitly set to False. Once the feature
+    is stabilized, the default will be changed to mutable=False.
+    """
+    return _CONFIG_DEFINITIONS.get(key, {}).get("mutable", True)
+
+
+def _get_config_dict(
+    config: LMCacheEngineConfig,
+    keys: Union[list[str], Optional[dict[str, dict[str, Any]]]] = None,
+) -> Dict[str, Any]:
+    """Get the config dict filtered by keys"""
+    if keys is None:
+        keys = _CONFIG_DEFINITIONS
+    return {key: getattr(config, key) for key in keys if hasattr(config, key)}
+
+
+@router.get("/conf")
+async def get_config(request: Request, names: Optional[str] = None):
+    config = request.app.state.lmcache_adapter.config
+    # Parse query parameter names (comma-separated list of config names)
+    keys = names.split(",") if names else None
+    config_dict = _get_config_dict(config, keys)
+    return PlainTextResponse(
+        content=json.dumps(config_dict, indent=2), media_type="text/plain"
+    )
+
+
+@router.get("/meta")
+async def get_metadata(request: Request, names: Optional[str] = None):
+    """
+    Get metadata of the cache engine
+    """
+    metadata = request.app.state.lmcache_adapter.lmcache_engine_metadata
+
+    if names:
+        attr_list = names.split(",")
+        metadata_dict = {}
+        for attr in attr_list:
+            if (
+                hasattr(metadata, attr)
+                and not attr.startswith("__")
+                and not callable(getattr(metadata, attr))
+            ):
+                value = getattr(metadata, attr)
+                if isinstance(value, (torch.dtype, torch.Size)):
+                    value = str(value)
+                metadata_dict[attr] = value
+    else:
+        metadata_dict = {
+            attr: getattr(metadata, attr)
+            for attr in dir(metadata)
+            if not attr.startswith("__") and not callable(getattr(metadata, attr))
+        }
+        for key, value in metadata_dict.items():
+            if isinstance(value, (torch.dtype, torch.Size)):
+                metadata_dict[key] = str(value)
+
+    return PlainTextResponse(
+        content=json.dumps(metadata_dict, indent=2, default=str),
+        media_type="text/plain",
+    )
+
+
+@router.post("/conf")
+async def set_config(request: Request):
+    """Set config values dynamically.
+
+    NOTE: Currently experimental — all configs are mutable at runtime
+    by default unless explicitly set "mutable": False in
+    _CONFIG_DEFINITIONS. The default will change to immutable once
+    the feature is stabilized.
+
+    Request body should be JSON with config name-value pairs:
+        {"min_retrieve_tokens": 512, "save_decode_cache": true}
+
+    Returns:
+        PlainTextResponse: JSON response with updated config values
+    """
+    config = request.app.state.lmcache_adapter.config
+
+    try:
+        body = await request.json()
+    except Exception as e:
+        return JSONResponse(
+            content={"error": "Invalid JSON body", "message": str(e)},
+            status_code=400,
+        )
+
+    if not isinstance(body, dict):
+        return JSONResponse(
+            content={"error": "Request body must be a JSON object"},
+            status_code=400,
+        )
+
+    updated: Dict[str, Any] = {}
+    errors: Dict[str, str] = {}
+
+    for key, value in body.items():
+        if not _is_mutable_config(key):
+            errors[key] = "Config is not mutable at runtime"
+            continue
+
+        if key not in _CONFIG_DEFINITIONS:
+            errors[key] = "Unknown config"
+            continue
+
+        # Use shared validation and setting function
+        if validate_and_set_config_value(config, key, value):
+            updated[key] = getattr(config, key)
+            logger.info("Config %s updated to %s", key, updated[key])
+        else:
+            errors[key] = "Failed to set config value"
+
+    result: Dict[str, Any] = {"updated": updated}
+    if errors:
+        result["errors"] = errors
+
+    status_code = 400 if errors and not updated else 200
+    return JSONResponse(
+        content=result,
+        status_code=status_code,
+    )

lmcache/v1/internal_api_server/vllm/freeze_api.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import json
+
+# Third Party
+from fastapi import APIRouter
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+
+router = APIRouter()
+
+
+@router.put("/freeze/enable")
+async def enable_freeze(request: Request):
+    """
+    Enable freeze mode for the LMCache engine.
+
+    When freeze mode is enabled:
+    - All store operations will be skipped (no new data stored)
+    - Only local_cpu backend will be used for retrieval
+    - No admit/evict messages will be generated
+    This protects the local_cpu hot cache from changes.
+
+    Args:
+        request (Request): The FastAPI request object containing application state.
+
+    Returns:
+        PlainTextResponse: A JSON response indicating the operation status.
+
+    Example:
+        ```bash
+        curl -X PUT "http://localhost:8000/freeze/enable"
+        # Response: {"status": "success", "freeze": true}
+        ```
+    """
+    try:
+        lmcache_adapter = request.app.state.lmcache_adapter
+        lmcache_engine = getattr(lmcache_adapter, "lmcache_engine", None)
+        if not lmcache_engine:
+            error_info = {
+                "error": "/freeze/enable API is unavailable",
+                "message": "LMCache engine not configured.",
+            }
+            return PlainTextResponse(
+                content=json.dumps(error_info, indent=2),
+                media_type="application/json",
+                status_code=503,  # Service Unavailable
+            )
+
+        lmcache_engine.freeze(True)
+        success_info = {
+            "status": "success",
+            "freeze": True,
+            "message": "Freeze mode enabled successfully",
+        }
+        return PlainTextResponse(
+            content=json.dumps(success_info, indent=2),
+            media_type="application/json",
+        )
+    except Exception as e:
+        error_msg = "Failed to enable freeze mode"
+        error_info = {"error": error_msg, "message": str(e)}
+        return PlainTextResponse(
+            content=json.dumps(error_info, indent=2),
+            media_type="application/json",
+            status_code=500,
+        )
+
+
+@router.put("/freeze/disable")
+async def disable_freeze(request: Request):
+    """
+    Disable freeze mode for the LMCache engine.
+
+    When freeze mode is disabled, store operations will proceed normally.
+
+    Args:
+        request (Request): The FastAPI request object containing application state.
+
+    Returns:
+        PlainTextResponse: A JSON response indicating the operation status.
+
+    Example:
+        ```bash
+        curl -X PUT "http://localhost:8000/freeze/disable"
+        # Response: {"status": "success", "freeze": false}
+        ```
+    """
+    try:
+        lmcache_adapter = request.app.state.lmcache_adapter
+        lmcache_engine = getattr(lmcache_adapter, "lmcache_engine", None)
+        if not lmcache_engine:
+            error_info = {
+                "error": "/freeze/disable API is unavailable",
+                "message": "LMCache engine not configured.",
+            }
+            return PlainTextResponse(
+                content=json.dumps(error_info, indent=2),
+                media_type="application/json",
+                status_code=503,  # Service Unavailable
+            )
+
+        lmcache_engine.freeze(False)
+        success_info = {
+            "status": "success",
+            "freeze": False,
+            "message": "Freeze mode disabled successfully",
+        }
+        return PlainTextResponse(
+            content=json.dumps(success_info, indent=2),
+            media_type="application/json",
+        )
+    except Exception as e:
+        error_msg = "Failed to disable freeze mode"
+        error_info = {"error": error_msg, "message": str(e)}
+        return PlainTextResponse(
+            content=json.dumps(error_info, indent=2),
+            media_type="application/json",
+            status_code=500,
+        )
+
+
+@router.get("/freeze/status")
+async def get_freeze_status(request: Request):
+    """
+    Get the current freeze mode status of the LMCache engine.
+
+    Args:
+        request (Request): The FastAPI request object containing application state.
+
+    Returns:
+        PlainTextResponse: JSON response with current freeze mode status.
+
+    Example:
+        ```bash
+        curl -X GET "http://localhost:8000/freeze/status"
+        # Response: {"status": "success", "freeze": true}
+        ```
+    """
+    try:
+        lmcache_adapter = request.app.state.lmcache_adapter
+        lmcache_engine = getattr(lmcache_adapter, "lmcache_engine", None)
+        if not lmcache_engine:
+            error_info = {
+                "error": "/freeze/status API is unavailable",
+                "message": "LMCache engine not configured.",
+            }
+            return PlainTextResponse(
+                content=json.dumps(error_info, indent=2),
+                media_type="application/json",
+                status_code=503,  # Service Unavailable
+            )
+
+        freeze_mode = lmcache_engine.is_frozen()
+        mode_str = "enabled" if freeze_mode else "disabled"
+        success_info = {
+            "status": "success",
+            "freeze": freeze_mode,
+            "message": "Freeze mode is " + mode_str,
+        }
+        return PlainTextResponse(
+            content=json.dumps(success_info, indent=2),
+            media_type="application/json",
+        )
+    except Exception as e:
+        error_msg = "Failed to get freeze mode status"
+        error_info = {"error": error_msg, "message": str(e)}
+        return PlainTextResponse(
+            content=json.dumps(error_info, indent=2),
+            media_type="application/json",
+            status_code=500,
+        )

lmcache/v1/internal_api_server/vllm/hot_cache_api.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import json
+
+# Third Party
+from fastapi import APIRouter
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+
+router = APIRouter()
+
+
+def _get_engine(request: Request):
+    """
+    Extract the LMCache engine from the request state.
+
+    Returns:
+        A tuple of (engine, error_response). If engine is None,
+        error_response contains a PlainTextResponse to return.
+    """
+    adapter = request.app.state.lmcache_adapter
+    engine = getattr(adapter, "lmcache_engine", None)
+    if not engine:
+        error_info = {
+            "error": "hot_cache API is unavailable",
+            "message": "LMCache engine not configured.",
+        }
+        return None, PlainTextResponse(
+            content=json.dumps(error_info, indent=2),
+            media_type="application/json",
+            status_code=503,
+        )
+    return engine, None
+
+
+@router.put("/hot_cache/enable")
+async def enable_hot_cache(request: Request):
+    """
+    Enable hot cache for the LocalCPUBackend.
+
+    When hot cache is enabled, frequently accessed KV cache
+    data will be kept in CPU memory for faster retrieval.
+
+    Args:
+        request: The FastAPI request object.
+
+    Returns:
+        PlainTextResponse: JSON response with operation status.
+
+    Example:
+        ```bash
+        curl -X PUT "http://localhost:8000/hot_cache/enable"
+        # Response: {"status": "success", "hot_cache": true}
+        ```
+    """
+    try:
+        engine, err = _get_engine(request)
+        if err:
+            return err
+
+        engine.set_hot_cache(True)
+        return PlainTextResponse(
+            content=json.dumps(
+                {
+                    "status": "success",
+                    "hot_cache": True,
+                    "message": "Hot cache enabled successfully",
+                },
+                indent=2,
+            ),
+            media_type="application/json",
+        )
+    except Exception as e:
+        return PlainTextResponse(
+            content=json.dumps(
+                {
+                    "error": "Failed to enable hot cache",
+                    "message": str(e),
+                },
+                indent=2,
+            ),
+            media_type="application/json",
+            status_code=500,
+        )
+
+
+@router.put("/hot_cache/disable")
+async def disable_hot_cache(request: Request):
+    """
+    Disable hot cache for the LocalCPUBackend.
+
+    When hot cache is disabled, existing hot cache entries
+    will be cleared and no new data will be written.
+
+    Args:
+        request: The FastAPI request object.
+
+    Returns:
+        PlainTextResponse: JSON response with operation status.
+
+    Example:
+        ```bash
+        curl -X PUT "http://localhost:8000/hot_cache/disable"
+        # Response: {"status": "success", "hot_cache": false}
+        ```
+    """
+    try:
+        engine, err = _get_engine(request)
+        if err:
+            return err
+
+        engine.set_hot_cache(False)
+        return PlainTextResponse(
+            content=json.dumps(
+                {
+                    "status": "success",
+                    "hot_cache": False,
+                    "message": ("Hot cache disabled successfully"),
+                },
+                indent=2,
+            ),
+            media_type="application/json",
+        )
+    except Exception as e:
+        return PlainTextResponse(
+            content=json.dumps(
+                {
+                    "error": "Failed to disable hot cache",
+                    "message": str(e),
+                },
+                indent=2,
+            ),
+            media_type="application/json",
+            status_code=500,
+        )
+
+
+@router.get("/hot_cache/status")
+async def get_hot_cache_status(request: Request):
+    """
+    Get the current hot cache status of LocalCPUBackend.
+
+    Args:
+        request: The FastAPI request object.
+
+    Returns:
+        PlainTextResponse: JSON response with hot cache status.
+
+    Example:
+        ```bash
+        curl -X GET "http://localhost:8000/hot_cache/status"
+        # Response: {"status": "success", "hot_cache": true}
+        ```
+    """
+    try:
+        engine, err = _get_engine(request)
+        if err:
+            return err
+
+        enabled = engine.is_hot_cache_enabled()
+        mode_str = "enabled" if enabled else "disabled"
+        return PlainTextResponse(
+            content=json.dumps(
+                {
+                    "status": "success",
+                    "hot_cache": enabled,
+                    "message": "Hot cache is " + mode_str,
+                },
+                indent=2,
+            ),
+            media_type="application/json",
+        )
+    except Exception as e:
+        return PlainTextResponse(
+            content=json.dumps(
+                {
+                    "error": "Failed to get hot cache status",
+                    "message": str(e),
+                },
+                indent=2,
+            ),
+            media_type="application/json",
+            status_code=500,
+        )