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

lmcache/v1/api_server/__main__.py

@@ -0,0 +1,537 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from contextlib import asynccontextmanager
+from typing import Any, Dict, List, Optional, Tuple
+import argparse
+import asyncio
+import json
+import os
+import sys
+import uuid
+
+# Add project root to Python path for local development
+sys.path.insert(
+    0,
+    os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    ),
+)
+
+# Third Party
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+import uvicorn
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.config import (
+    load_controller_config_with_overrides,
+)
+from lmcache.v1.cache_controller.controller_manager import LMCacheControllerManager
+from lmcache.v1.cache_controller.message import (  # noqa: E501
+    CheckFinishMsg,
+    CheckFinishRetMsg,
+    ClearMsg,
+    ClearRetMsg,
+    CompressMsg,
+    CompressRetMsg,
+    DecompressMsg,
+    DecompressRetMsg,
+    ErrorMsg,
+    HealthMsg,
+    HealthRetMsg,
+    LookupMsg,
+    LookupRetMsg,
+    MoveMsg,
+    MoveRetMsg,
+    PinMsg,
+    PinRetMsg,
+    QueryInstMsg,
+    QueryInstRetMsg,
+    QueryWorkerInfoMsg,
+    QueryWorkerInfoRetMsg,
+    WorkerInfo,
+)
+from lmcache.v1.config_base import parse_command_line_extra_params
+from lmcache.v1.internal_api_server.api_registry import APIRegistry
+
+logger = init_logger(__name__)
+
+
+def parse_extra_params(extra_args: list) -> Dict[str, Any]:
+    """Parse extra parameters in key=value format"""
+    params = {}
+    for arg in extra_args:
+        if "=" in arg:
+            key, value = arg.split("=", 1)
+            key = key.lstrip("-")
+            try:
+                if value.lower() in ("true", "false"):
+                    params[key] = value.lower() == "true"
+                elif value.isdigit():
+                    params[key] = int(value)
+                elif value.replace(".", "", 1).isdigit():
+                    params[key] = float(value)
+                else:
+                    params[key] = value
+            except ValueError:
+                params[key] = value
+            logger.info(f"Extra parameter: {key} = {params[key]}")
+    return params
+
+
+def create_app(
+    controller_urls: dict[str, str],
+    health_check_interval: int,
+    lmcache_worker_timeout: int,
+) -> FastAPI:
+    """
+    Create a FastAPI application with endpoints for LMCache operations.
+    """
+    lmcache_controller_manager = LMCacheControllerManager(
+        controller_urls, health_check_interval, lmcache_worker_timeout
+    )
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        # Start background task here
+        lmcache_cluster_monitor_task = asyncio.create_task(
+            lmcache_controller_manager.start_all()
+        )
+        yield
+        # Optionally cancel the task on shutdown
+        lmcache_cluster_monitor_task.cancel()
+        try:
+            await lmcache_cluster_monitor_task
+        except asyncio.CancelledError:
+            pass
+
+    app = FastAPI(lifespan=lifespan)
+    app.state.lmcache_controller_manager = lmcache_controller_manager
+
+    # Register internal APIs (only common APIs, not vllm-specific ones)
+    registry = APIRegistry(app)
+    registry.register_all_apis(categories=["common", "controller"])
+
+    # Add static files for frontend
+    project_root = os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    )
+    static_dir = os.path.join(
+        project_root,
+        "lmcache",
+        "v1",
+        "cache_controller",
+        "frontend",
+        "static",
+    )
+    if os.path.exists(static_dir):
+        app.mount("/static", StaticFiles(directory=static_dir), name="static")
+        logger.info("Controller frontend static files mounted at /static")
+    else:
+        logger.warning("Controller frontend static directory not found: %s", static_dir)
+
+    @app.get("/", response_class=HTMLResponse)
+    async def serve_frontend():
+        """Serve the Controller frontend HTML page."""
+        index_path = os.path.join(static_dir, "index.html")
+        if os.path.exists(index_path):
+            with open(index_path, "r") as f:
+                html_content = f.read()
+            return HTMLResponse(content=html_content)
+        else:
+            return HTMLResponse(
+                content="<h1>Controller Frontend not found</h1>"
+                "<p>Please build the frontend first.</p>",
+                status_code=404,
+            )
+
+    class QueryInstRequest(BaseModel):
+        event_id: str
+        ip: str
+
+    class QueryInstResponse(BaseModel):
+        event_id: str
+        res: str  # the instance id
+
+    @app.post("/query_instance")
+    async def query_instance(req: QueryInstRequest):
+        try:
+            event_id = "QueryInst" + str(uuid.uuid4())
+            msg = QueryInstMsg(
+                event_id=event_id,
+                ip=req.ip,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, QueryInstRetMsg)
+            return QueryInstResponse(
+                event_id=ret_msg.event_id,
+                res=ret_msg.instance_id,
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class LookupRequest(BaseModel):
+        tokens: List[int]
+
+    class LookupResponse(BaseModel):
+        event_id: str
+        # a list of (instance_id, location, token_count)
+        layout_info: Dict[str, Tuple[str, int]]
+
+    @app.post("/lookup", response_model=LookupResponse)
+    async def lookup(req: LookupRequest):
+        try:
+            event_id = "Lookup" + str(uuid.uuid4())
+            msg = LookupMsg(
+                event_id=event_id,
+                tokens=req.tokens,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, LookupRetMsg)
+            return LookupResponse(
+                event_id=ret_msg.event_id, layout_info=ret_msg.layout_info
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class ClearRequest(BaseModel):
+        instance_id: str
+        location: str
+
+    class ClearResponse(BaseModel):
+        event_id: str
+        num_tokens: int
+
+    @app.post("/clear", response_model=ClearResponse)
+    async def clear(req: ClearRequest):
+        try:
+            event_id = "Clear" + str(uuid.uuid4())
+            msg = ClearMsg(
+                event_id=event_id,
+                instance_id=req.instance_id,
+                location=req.location,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, ClearRetMsg)
+            return ClearResponse(
+                event_id=ret_msg.event_id, num_tokens=ret_msg.num_tokens
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class PinRequest(BaseModel):
+        instance_id: str
+        location: str
+        tokens: list[int]
+
+    class PinResponse(BaseModel):
+        event_id: str
+        num_tokens: int
+
+    @app.post("/pin", response_model=PinResponse)
+    async def pin(req: PinRequest):
+        try:
+            event_id = "Pin" + str(uuid.uuid4())
+            msg = PinMsg(
+                event_id=event_id,
+                instance_id=req.instance_id,
+                location=req.location,
+                tokens=req.tokens,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, PinRetMsg)
+            return PinResponse(event_id=ret_msg.event_id, num_tokens=ret_msg.num_tokens)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class CompressRequest(BaseModel):
+        instance_id: str
+        method: str
+        location: str
+        tokens: Optional[List[int]] = []
+
+    class CompressResponse(BaseModel):
+        event_id: str
+        num_tokens: int
+
+    class DecompressRequest(BaseModel):
+        instance_id: str
+        method: str
+        location: str
+        tokens: Optional[List[int]] = []
+
+    class DecompressResponse(BaseModel):
+        event_id: str
+        num_tokens: int
+
+    @app.post("/compress", response_model=CompressResponse)
+    async def compress(req: CompressRequest):
+        try:
+            event_id = "Compress" + str(uuid.uuid4())
+            msg = CompressMsg(
+                event_id=event_id,
+                instance_id=req.instance_id,
+                method=req.method,
+                location=req.location,
+                tokens=req.tokens,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, CompressRetMsg)
+            return CompressResponse(
+                event_id=ret_msg.event_id, num_tokens=ret_msg.num_tokens
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    @app.post("/decompress", response_model=DecompressResponse)
+    async def decompress(req: DecompressRequest):
+        try:
+            event_id = "Decompress" + str(uuid.uuid4())
+            msg = DecompressMsg(
+                event_id=event_id,
+                instance_id=req.instance_id,
+                method=req.method,
+                location=req.location,
+                tokens=req.tokens,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert isinstance(ret_msg, DecompressRetMsg)
+            return DecompressResponse(
+                event_id=ret_msg.event_id, num_tokens=ret_msg.num_tokens
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class MoveRequest(BaseModel):
+        # (instance_id, location)
+        old_position: Tuple[str, str]
+        new_position: Tuple[str, str]
+        tokens: Optional[List[int]] = []
+        should_copy: Optional[bool] = False
+
+    class MoveResponse(BaseModel):
+        event_id: str
+        num_tokens: int
+
+    @app.post("/move", response_model=MoveResponse)
+    async def move(req: MoveRequest):
+        try:
+            event_id = "Move" + str(uuid.uuid4())
+            msg = MoveMsg(
+                event_id=event_id,
+                old_position=req.old_position,
+                new_position=req.new_position,
+                tokens=req.tokens,
+                copy=req.should_copy,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, MoveRetMsg)
+            return MoveResponse(
+                event_id=ret_msg.event_id,
+                num_tokens=ret_msg.num_tokens,
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class HealthRequest(BaseModel):
+        instance_id: str
+
+    class HealthResponse(BaseModel):
+        event_id: str
+        # worker_id -> error_code
+        error_codes: dict[int, int]
+
+    @app.post("/health", response_model=HealthResponse)
+    async def health(req: HealthRequest):
+        try:
+            event_id = "health" + str(uuid.uuid4())
+            msg = HealthMsg(
+                event_id=event_id,
+                instance_id=req.instance_id,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, HealthRetMsg)
+            return HealthResponse(
+                event_id=ret_msg.event_id, error_codes=ret_msg.error_codes
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class CheckFinishRequest(BaseModel):
+        event_id: str
+
+    class CheckFinishResponse(BaseModel):
+        status: str
+
+    @app.post("/check_finish", response_model=CheckFinishResponse)
+    async def check_finish(req: CheckFinishRequest):
+        try:
+            msg = CheckFinishMsg(
+                event_id=req.event_id,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, CheckFinishRetMsg)
+            return CheckFinishResponse(status=ret_msg.status)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    class QueryWorkerInfoRequest(BaseModel):
+        instance_id: str
+        worker_ids: Optional[list[int]] = None
+
+    class QueryWorkerInfoResponse(BaseModel):
+        event_id: str
+        worker_infos: list[WorkerInfo]
+
+    @app.post("/query_worker_info", response_model=QueryWorkerInfoResponse)
+    async def query_worker_info(req: QueryWorkerInfoRequest):
+        try:
+            event_id = "QueryWorkerInfo" + str(uuid.uuid4())
+            msg = QueryWorkerInfoMsg(
+                event_id=event_id,
+                instance_id=req.instance_id,
+                worker_ids=req.worker_ids,
+            )
+            ret_msg = await lmcache_controller_manager.handle_orchestration_message(msg)
+            assert not isinstance(ret_msg, ErrorMsg), ret_msg.error
+            assert isinstance(ret_msg, QueryWorkerInfoRetMsg)
+            return QueryWorkerInfoResponse(
+                event_id=ret_msg.event_id, worker_infos=ret_msg.worker_infos
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    return app
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config", type=str, help="Path to controller configuration file"
+    )
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=9000)
+    parser.add_argument(
+        "--monitor-ports",
+        type=json.loads,
+        default=None,
+        help='JSON string of monitor ports, e.g. \'{"pull": 8300, "reply": 8400}\'',
+    )
+    parser.add_argument(
+        "--monitor-port",
+        type=int,
+        default=9001,
+        help="The controller pull port to maintain backward compatibility.",
+    )
+    parser.add_argument(
+        "--health-check-interval",
+        type=int,
+        default=-1,
+        help="Health check interval in secs, default is -1, which means disabled.",
+    )
+    parser.add_argument(
+        "--lmcache-worker-timeout",
+        type=int,
+        default=300,
+        help="The lmcache worker timeout in seconds.",
+    )
+
+    # Parse known args first, then handle extra parameters
+    args, extra = parser.parse_known_args()
+    extra_params = parse_command_line_extra_params(extra)
+
+    try:
+        # Build overrides dictionary from command-line arguments
+        override_dict = {}
+
+        # Map command-line arguments to config keys
+        arg_mappings = {
+            "host": "controller_host",
+            "port": "controller_port",
+            "monitor_ports": "controller_monitor_ports",
+            "health_check_interval": "health_check_interval",
+            "lmcache_worker_timeout": "lmcache_worker_timeout",
+        }
+
+        for arg_name, config_key in arg_mappings.items():
+            arg_value = getattr(args, arg_name)
+            if arg_value is not None:
+                override_dict[config_key] = arg_value
+
+        # Add extra parameters
+        if extra_params:
+            override_dict.update(extra_params)
+
+        # Load configuration using the generic utility function
+        # This replaces the previous manual config loading code
+        config = load_controller_config_with_overrides(
+            config_file_path=args.config,
+            overrides=override_dict,
+        )
+
+        # Build controller URLs from config or arguments
+        if config.controller_monitor_ports is not None:
+            controller_urls = {
+                "pull": (
+                    f"{config.controller_host}:"
+                    f"{config.controller_monitor_ports['pull']}"
+                ),
+                "reply": (
+                    f"{config.controller_host}:"
+                    f"{config.controller_monitor_ports['reply']}"
+                ),
+                "heartbeat": (
+                    f"{config.controller_host}:"
+                    f"{config.controller_monitor_ports['heartbeat']}"
+                    if config.controller_monitor_ports.get("heartbeat")
+                    else None
+                ),
+            }
+        else:
+            if args.monitor_port != 9001:  # Only warn if explicitly set
+                logger.warning(
+                    "Argument --monitor-port will be deprecated soon. "
+                    "Please use --monitor-ports instead."
+                )
+            controller_urls = {
+                "pull": f"{config.controller_host}:{args.monitor_port}",
+                "reply": None,
+                "heartbeat": None,
+            }
+
+        # Use config values for health check and timeout
+        health_check_interval = config.health_check_interval
+        lmcache_worker_timeout = config.lmcache_worker_timeout
+
+        app = create_app(controller_urls, health_check_interval, lmcache_worker_timeout)
+
+        logger.info(
+            f"Starting LMCache controller at "
+            f"{config.controller_host}:{config.controller_port}"
+        )
+        ports_message = f"Monitoring lmcache workers at ports {controller_urls}"
+        logger.info(ports_message)
+        logger.info(f"Health check interval: {health_check_interval}s")
+        logger.info(f"Worker timeout: {lmcache_worker_timeout}s")
+
+        uvicorn.run(app, host=config.controller_host, port=config.controller_port)
+    except TimeoutError as e:
+        logger.error(e)
+    except Exception as e:
+        logger.error(f"Failed to start controller: {e}", exc_info=True)
+        sys.exit(1)  # Exit with error code
+
+
+if __name__ == "__main__":
+    main()

lmcache/v1/basic_check.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import argparse
+import asyncio
+
+# First Party
+from lmcache.v1.check import registry
+
+model_name = "/lmcache_test_model/"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="LMCache basic check Tool")
+    parser.add_argument(
+        "--mode",
+        required=True,
+        help="Operation mode (e.g. test_remote, test_storage_manager). "
+        "Use 'list' to show available modes",
+    )
+    parser.add_argument("--model", default=model_name, help="model name")
+    parser.add_argument(
+        "--num-keys",
+        type=int,
+        default=5,
+        help="Number of keys for gen mode or test iterations "
+        "for test_* modes (default: 5)",
+    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=16,
+        help="Concurrency level for generation (gen mode only)",
+    )
+    parser.add_argument(
+        "--offset",
+        type=int,
+        default=0,
+        help="Offset for key generation (gen mode only)",
+    )
+    parser.add_argument(
+        "--l2-adapter",
+        dest="l2_adapter",
+        action="append",
+        default=[],
+        type=str,
+        metavar="JSON",
+        help="L2 adapter spec as JSON (test_l2_adapter mode). "
+        'e.g. \'{"type":"mock","max_size_gb":1}\'.',
+    )
+    parser.add_argument(
+        "--obj-size",
+        dest="obj_size",
+        type=int,
+        default=None,
+        help="Object size in number of elements (default: 1024)",
+    )
+    parser.add_argument(
+        "--kv-dtype",
+        dest="kv_dtype",
+        type=str,
+        default=None,
+        help="KV dtype, e.g. float32, bfloat16, float16 (default depends on mode)",
+    )
+    parser.add_argument(
+        "--settle-time",
+        dest="settle_time",
+        type=float,
+        default=0.0,
+        help="Seconds to wait after store before load "
+        "(default: 0, useful for remote backends)",
+    )
+    return parser.parse_args()
+
+
+async def main():
+    args = parse_args()
+
+    # List available modes if requested
+    if args.mode == "list":
+        registry.load_modes()
+        print("Available check modes:")
+        for mode_name in registry.modes:
+            print(f"  - {mode_name}")
+        return
+
+    # Get the requested mode function
+    mode_func = registry.get_mode(args.mode)
+    if not mode_func:
+        print(
+            f"Error: Unknown mode '{args.mode}'. "
+            "Use '--mode list' to see available modes."
+        )
+        return
+
+    # Prepare arguments for the mode function
+    mode_args = {
+        "model": args.model,
+        "num_keys": args.num_keys,
+        "concurrency": args.concurrency,
+        "offset": args.offset,
+        "l2_adapter": args.l2_adapter,
+        "obj_size": args.obj_size,
+        "kv_dtype": args.kv_dtype,
+        "settle_time": args.settle_time,
+    }
+
+    # Execute the mode function
+    await mode_func(**mode_args)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

lmcache/v1/cache_controller/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.cache_controller.executor import LMCacheClusterExecutor  # noqa: E501
+from lmcache.v1.cache_controller.worker import LMCacheWorker  # noqa: E501
+
+__all__ = [
+    "LMCacheClusterExecutor",
+    "LMCacheWorker",
+]

lmcache/v1/cache_controller/commands/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Heartbeat commands module
+
+This module provides the command abstraction for the heartbeat mechanism.
+Commands can be sent from controller to workers through heartbeat responses.
+"""
+
+# First Party
+from lmcache.v1.cache_controller.commands.base import HeartbeatCommand
+from lmcache.v1.cache_controller.commands.full_sync import FullSyncCommand
+
+__all__ = [
+    "HeartbeatCommand",
+    "FullSyncCommand",
+]