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

lmcache/v1/multiprocess/http_server.py

@@ -0,0 +1,235 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from contextlib import asynccontextmanager
+import argparse
+
+# Third Party
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+import torch
+import uvicorn
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.distributed.config import (
+    StorageManagerConfig,
+    add_storage_manager_args,
+    parse_args_to_config,
+)
+from lmcache.v1.mp_observability.config import (
+    ObservabilityConfig,
+    add_observability_args,
+    parse_args_to_observability_config,
+)
+from lmcache.v1.mp_observability.event_bus import get_event_bus
+from lmcache.v1.multiprocess.config import (
+    HTTPFrontendConfig,
+    MPServerConfig,
+    add_http_frontend_args,
+    add_mp_server_args,
+    parse_args_to_http_frontend_config,
+    parse_args_to_mp_server_config,
+)
+from lmcache.v1.multiprocess.mp_runtime_plugin_launcher import (
+    MPRuntimePluginLauncher,
+)
+
+logger = init_logger(__name__)
+
+
+# Module-level config holders, set by run_http_server() before FastAPI startup.
+# Stored in a dict so the lifespan closure captures the mutable container.
+_configs: dict = {}
+
+
+# ----------------------------
+# FastAPI lifespan for initialization and cleanup
+# ----------------------------
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """
+    Manage the lifecycle of the LMCache HTTP server.
+
+    On startup: Initialize ZMQ server and cache engine.
+    On shutdown: Clean up ZMQ server resources.
+    """
+    # Startup
+    logger.info(
+        "Starting LMCache HTTP server... (CUDA available: %s)",
+        torch.cuda.is_available(),
+    )
+    mp_config = _configs["mp"]
+    if mp_config.engine_type == "blend":
+        # First Party
+        from lmcache.v1.multiprocess.blend_server_v2 import run_cache_server
+    else:
+        # First Party
+        from lmcache.v1.multiprocess.server import run_cache_server
+
+    result = run_cache_server(
+        mp_config=mp_config,
+        storage_manager_config=_configs["storage_manager"],
+        obs_config=_configs["observability"],
+        return_engine=True,
+    )
+    assert result is not None, "run_cache_server returned None with return_engine=True"
+    zmq_server, engine = result
+
+    # Launch runtime plugins if configured. Plugins receive the full
+    # server config (including HTTP host/port) via the
+    # LMCACHE_RUNTIME_PLUGIN_CONFIG environment variable.
+    plugin_launcher = None
+    if mp_config.runtime_plugin_config.locations:
+        extra_kwargs = {}
+        http_config = _configs.get("http")
+        if http_config is not None:
+            extra_kwargs["http_config"] = http_config
+        plugin_launcher = MPRuntimePluginLauncher(
+            runtime_plugin_config=mp_config.runtime_plugin_config,
+            mp_config=mp_config,
+            storage_manager_config=_configs["storage_manager"],
+            obs_config=_configs["observability"],
+            **extra_kwargs,
+        )
+        plugin_launcher.launch_plugins()
+
+    app.state.zmq_server = zmq_server
+    app.state.engine = engine
+    app.state.plugin_launcher = plugin_launcher
+    logger.info("LMCache HTTP server initialized")
+
+    yield
+
+    # Shutdown
+    logger.info("Shutting down LMCache HTTP server...")
+    launcher = getattr(app.state, "plugin_launcher", None)
+    if launcher is not None:
+        launcher.stop_plugins()
+    get_event_bus().stop()
+    if hasattr(app.state, "zmq_server") and app.state.zmq_server is not None:
+        app.state.zmq_server.close()
+    logger.info("LMCache HTTP server stopped")
+
+
+app = FastAPI(title="LMCache HTTP API", version="1.0.0", lifespan=lifespan)
+
+
+@app.get("/")
+async def root():
+    return {"status": "ok", "service": "LMCache HTTP API"}
+
+
+@app.get("/api/healthcheck")
+async def healthcheck(request: Request):
+    """
+    Health check endpoint for k8s liveness/readiness probes.
+
+    Checks:
+        - HTTP server is alive (implicit: if you get a response)
+        - Cache engine is alive
+    """
+    engine = getattr(request.app.state, "engine", None)
+    if engine is None:
+        return JSONResponse(
+            status_code=503,
+            content={"status": "unhealthy", "reason": "engine not initialized"},
+        )
+
+    return {"status": "healthy"}
+
+
+@app.post("/api/clear-cache")
+async def clear_cache(request: Request):
+    """
+    Force-clear all KV cache data stored in L1 (CPU) memory.
+
+    This clears all objects including those with active read/write locks.
+    In-flight store or prefetch operations may be corrupted.
+    """
+    engine = getattr(request.app.state, "engine", None)
+    if engine is None:
+        return JSONResponse(
+            status_code=503,
+            content={"status": "error", "reason": "engine not initialized"},
+        )
+
+    engine.clear()
+    logger.info("Cache cleared via HTTP API")
+    return {"status": "ok"}
+
+
+@app.get("/api/status")
+async def status(request: Request):
+    """
+    Detailed status endpoint for inspecting internal state of all
+    MP components (L1 cache, L2 adapters, controllers, sessions).
+    """
+    engine = getattr(request.app.state, "engine", None)
+    if engine is None:
+        return JSONResponse(
+            status_code=503,
+            content={"error": "engine not initialized"},
+        )
+    return engine.report_status()
+
+
+def run_http_server(
+    http_config: HTTPFrontendConfig,
+    mp_config: MPServerConfig,
+    storage_manager_config: StorageManagerConfig,
+    obs_config: ObservabilityConfig,
+) -> None:
+    """
+    Run the LMCache HTTP server with integrated MP (ZMQ) server.
+
+    Args:
+        http_config: Configuration for the HTTP frontend
+        mp_config: Configuration for the ZMQ multiprocess server
+        storage_manager_config: Configuration for the storage manager
+        obs_config: Configuration for the observability stack
+    """
+    _configs["mp"] = mp_config
+    _configs["storage_manager"] = storage_manager_config
+    _configs["observability"] = obs_config
+    _configs["http"] = http_config
+
+    config = uvicorn.Config(
+        app=app,
+        host=http_config.http_host,
+        port=http_config.http_port,
+        log_level="info",
+        access_log=True,
+    )
+    server = uvicorn.Server(config)
+
+    logger.info(
+        "Starting LMCache HTTP server on http://%s:%d",
+        http_config.http_host,
+        http_config.http_port,
+    )
+    server.run()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="LMCache HTTP Server with integrated MP Cache Server"
+    )
+    add_http_frontend_args(parser)
+    add_mp_server_args(parser)
+    add_storage_manager_args(parser)
+    add_observability_args(parser)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    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)
+    run_http_server(
+        http_config=http_config,
+        mp_config=mp_config,
+        storage_manager_config=storage_manager_config,
+        obs_config=obs_config,
+    )

lmcache/v1/multiprocess/mp_runtime_plugin_launcher.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Runtime plugin launcher for multiprocess (MP) mode.
+
+Unlike the non-MP RuntimePluginLauncher which receives a single
+LMCacheEngineConfig, the MP mode has multiple independent config
+dataclasses (MPServerConfig, StorageManagerConfig,
+ObservabilityConfig, etc.).  This launcher aggregates them all
+into a single JSON blob so plugins get the full server config
+via the LMCACHE_RUNTIME_PLUGIN_CONFIG environment variable.
+"""
+
+# Standard
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Any
+import json
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.plugin.runtime_plugin_launcher import (
+    RuntimePluginLauncher,
+)
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.multiprocess.config import (
+        RuntimePluginConfig,
+    )
+
+logger = init_logger(__name__)
+
+
+def _safe_asdict(obj: Any) -> dict[str, Any]:
+    """Convert a dataclass to dict, falling back to str()
+    for non-serializable fields."""
+    # Verify it's a dataclass instance at runtime
+    if not hasattr(obj, "__dataclass_fields__"):
+        raise TypeError("Expected a dataclass instance, got %s" % type(obj).__name__)
+    raw = asdict(obj)
+    return _make_json_safe(raw)  # type: ignore[return-value]
+
+
+def _make_json_safe(obj: Any) -> Any:
+    """Recursively ensure all values are JSON-serializable."""
+    if isinstance(obj, dict):
+        return {k: _make_json_safe(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_make_json_safe(v) for v in obj]
+    if isinstance(obj, (str, int, float, bool, type(None))):
+        return obj
+    # Fallback: convert to string representation
+    return str(obj)
+
+
+@dataclass
+class _MPPluginConfig:
+    """Thin wrapper that satisfies RuntimePluginLauncher's
+    config duck-type contract (runtime_plugin_locations +
+    to_json)."""
+
+    runtime_plugin_locations: list[str]
+    extra_config: dict[str, Any]
+    configs_dict: dict[str, Any]
+
+    def to_json(self) -> str:
+        merged = dict(self.configs_dict)
+        if self.extra_config:
+            merged["runtime_plugin_extra_config"] = _make_json_safe(self.extra_config)
+        return json.dumps(merged)
+
+
+class MPRuntimePluginLauncher:
+    """Launch runtime plugins in MP mode with all server
+    configs serialized into the environment.
+
+    Usage::
+
+        launcher = MPRuntimePluginLauncher(
+            runtime_plugin_config=runtime_plugin_config,
+            mp_config=mp_config,
+            storage_manager_config=sm_config,
+            obs_config=obs_config,
+        )
+        launcher.launch_plugins()
+        # ... on shutdown ...
+        launcher.stop_plugins()
+    """
+
+    def __init__(
+        self,
+        runtime_plugin_config: "RuntimePluginConfig",
+        **configs: object,
+    ) -> None:
+        """Initialize the MP runtime plugin launcher.
+
+        Aggregates arbitrary dataclass configs into a single
+        JSON blob and delegates to RuntimePluginLauncher.
+
+        Args:
+            runtime_plugin_config: RuntimePluginConfig with
+                locations and extra_config fields.
+            **configs: Dataclass config objects to serialize
+                and pass to plugins via environment variable.
+        """
+        # Build the aggregated JSON dict from all configs
+        aggregated: dict = {}
+        for name, cfg in configs.items():
+            aggregated[name] = _safe_asdict(cfg)
+
+        wrapper = _MPPluginConfig(
+            runtime_plugin_locations=runtime_plugin_config.locations,
+            extra_config=getattr(runtime_plugin_config, "extra_config", {}),
+            configs_dict=aggregated,
+        )
+        self._inner = RuntimePluginLauncher(
+            config=wrapper,
+            role=None,
+            worker_count=1,
+            worker_id=0,
+        )
+        logger.info(f"MPRuntimePluginLauncher initialized with {wrapper}")
+
+    def launch_plugins(self) -> None:
+        """Launch all configured plugins."""
+        self._inner.launch_plugins()
+
+    def stop_plugins(self) -> None:
+        """Terminate all plugin processes."""
+        self._inner.stop_plugins()