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

lmcache/v1/internal_api_server/common/periodic_thread_api.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+API endpoint for monitoring periodic threads.
+"""
+
+# Standard
+from typing import Optional
+
+# Third Party
+from fastapi import APIRouter, Query
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+# First Party
+from lmcache.v1.periodic_thread import PeriodicThreadRegistry, ThreadLevel
+
+router = APIRouter()
+
+
+@router.get("/periodic-threads")
+async def get_periodic_threads(
+    request: Request,
+    level: Optional[str] = Query(
+        None,
+        description="Filter by thread level (critical, high, medium, low)",
+    ),
+    running_only: bool = Query(
+        False,
+        description="Only show running threads",
+    ),
+    active_only: bool = Query(
+        False,
+        description="Only show active threads",
+    ),
+):
+    """
+    Get information about registered periodic threads.
+
+    Returns a summary of all periodic threads including:
+    - Total, running, and active counts by level
+    - Individual thread status with last run time and summary
+    """
+    registry = PeriodicThreadRegistry.get_instance()
+
+    # Get all threads
+    if level:
+        try:
+            thread_level = ThreadLevel(level.lower())
+            threads = registry.get_by_level(thread_level)
+        except ValueError:
+            return JSONResponse(
+                status_code=400,
+                content={
+                    "error": f"Invalid level: {level}. "
+                    f"Valid values: critical, high, medium, low"
+                },
+            )
+    else:
+        threads = registry.get_all()
+
+    # Apply filters
+    if running_only:
+        threads = [t for t in threads if t.is_running]
+    if active_only:
+        threads = [t for t in threads if t.is_active]
+
+    # Build response
+    thread_statuses = [t.get_status() for t in threads]
+
+    # Get summary
+    summary = registry.get_summary()
+
+    return JSONResponse(
+        content={
+            "summary": {
+                "total_count": summary["total_count"],
+                "running_count": summary["running_count"],
+                "active_count": summary["active_count"],
+                "by_level": summary["by_level"],
+            },
+            "threads": thread_statuses,
+        }
+    )
+
+
+@router.get("/periodic-threads/{thread_name}")
+async def get_periodic_thread(
+    request: Request,
+    thread_name: str,
+):
+    """
+    Get detailed information about a specific periodic thread.
+    """
+    registry = PeriodicThreadRegistry.get_instance()
+    thread = registry.get(thread_name)
+
+    if thread is None:
+        return JSONResponse(
+            status_code=404,
+            content={"error": f"Thread not found: {thread_name}"},
+        )
+
+    return JSONResponse(content=thread.get_status())
+
+
+@router.get("/periodic-threads-health")
+async def get_periodic_threads_health(request: Request):
+    """
+    Quick health check for periodic threads.
+
+    Returns:
+    - healthy: True if all critical/high level threads are active
+    - unhealthy_threads: List of inactive critical/high threads
+    """
+    registry = PeriodicThreadRegistry.get_instance()
+
+    unhealthy_threads = []
+
+    # Check critical and high level threads
+    for level in [ThreadLevel.CRITICAL, ThreadLevel.HIGH]:
+        for thread in registry.get_by_level(level):
+            if thread.is_running and not thread.is_active:
+                unhealthy_threads.append(
+                    {
+                        "name": thread.name,
+                        "level": thread.level.value,
+                        "last_run_ago": thread.get_status().get("last_run_ago"),
+                        "interval": thread.interval,
+                    }
+                )
+
+    return JSONResponse(
+        content={
+            "healthy": len(unhealthy_threads) == 0,
+            "unhealthy_count": len(unhealthy_threads),
+            "unhealthy_threads": unhealthy_threads,
+        }
+    )

lmcache/v1/internal_api_server/common/run_script_api.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Any
+import importlib
+
+# Third Party
+from fastapi import APIRouter
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+router = APIRouter()
+
+
+@router.post("/run_script")
+async def run_script(request: Request):
+    form_data = await request.form()
+    script_file = form_data.get("script")
+
+    if not script_file or not hasattr(script_file, "file"):
+        return PlainTextResponse("No script file provided", status_code=400)
+
+    script_content = await script_file.read()
+
+    try:
+        # Get allowed imports from config
+        config = request.app.state.lmcache_adapter.config
+        allowed_imports = config.script_allowed_imports or []
+
+        # Pre-import allowed modules
+        allowed_modules = {}
+        for module_name in allowed_imports:
+            try:
+                module = importlib.import_module(module_name)
+                allowed_modules[module_name] = module
+                logger.info(f"Imported allowed module: {module_name}")
+            except ImportError as e:
+                logger.warning(f"Failed to import module {module_name}: {e}")
+
+        # Create custom __import__ function that only allows configured modules
+        def restricted_import(name, globals=None, locals=None, fromlist=(), level=0):
+            if name in allowed_modules:
+                return allowed_modules[name]
+            raise ImportError(f"Import of '{name}' is not allowed")
+
+        restricted_globals = {
+            "__builtins__": {
+                "print": print,
+                "str": str,
+                "int": int,
+                "float": float,
+                "list": list,
+                "dict": dict,
+                "tuple": tuple,
+                "set": set,
+                "__import__": restricted_import,
+            },
+            "app": request.app,
+        }
+
+        restricted_locals: dict[str, Any] = {}
+
+        exec(script_content, restricted_globals, restricted_locals)
+
+        result = restricted_locals.get("result", "Script executed successfully")
+        return PlainTextResponse(str(result), media_type="text/plain")
+
+    except Exception as e:
+        return PlainTextResponse(f"Error executing script: {str(e)}", status_code=500)

lmcache/v1/internal_api_server/common/thread_api.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional
+import sys
+import threading
+import traceback
+
+# Third Party
+from fastapi import APIRouter, Query
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
+
+router = APIRouter()
+
+
+@router.get("/threads")
+async def get_threads(
+    request: Request,
+    name: Optional[str] = Query(
+        None, description="Filter by thread name (fuzzy match)"
+    ),
+    thread_id: Optional[int] = Query(None, description="Filter by thread ID"),
+):
+    """Return information about active threads with optional filtering"""
+    threads = threading.enumerate()
+
+    filtered_threads = []
+    for t in threads:
+        # Apply filters
+        if name and name.lower() not in t.name.lower():
+            continue
+        if thread_id and t.ident != thread_id:
+            continue
+        filtered_threads.append(t)
+
+    thread_info = []
+
+    for t in filtered_threads:
+        # Basic thread info with creation time
+        info = f"Thread: {t}\n"
+
+        # Get stack trace if available
+        try:
+            stack_frames = (
+                sys._current_frames().get(t.ident) if t.ident is not None else None
+            )
+            if stack_frames:
+                stack_trace = traceback.format_stack(stack_frames)
+                info += "Stack trace:\n" + "".join(stack_trace)
+            else:
+                info += "No stack trace available\n"
+        except AttributeError:
+            info += "Stack trace unavailable\n"
+
+        thread_info.append(info)
+
+    # Add summary section
+    summary = "\n\n=== Thread Summary ===\n"
+    summary += f"Total threads: {len(filtered_threads)}\n"
+
+    return PlainTextResponse(
+        content="\n\n".join(thread_info) + summary, media_type="text/plain"
+    )

lmcache/v1/internal_api_server/controller/__init__.py

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

lmcache/v1/internal_api_server/controller/key_stats_api.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import List
+
+# Third Party
+from fastapi import APIRouter, HTTPException, Request
+from pydantic import BaseModel
+
+router = APIRouter()
+
+
+class InstanceKeyStats(BaseModel):
+    instance_id: str
+    key_count: int
+    worker_count: int
+
+
+class KeyStatsResponse(BaseModel):
+    total_key_count: int
+    total_instance_count: int
+    total_worker_count: int
+    instances: List[InstanceKeyStats]
+
+
+@router.get("/controller/key-stats")
+async def get_key_stats(request: Request):
+    """
+    Get key statistics across all instances and workers.
+
+    Returns:
+        - Total key count across all instances
+        - Total instance count
+        - Total worker count
+        - Key count per instance
+    """
+    try:
+        controller_manager = getattr(
+            request.app.state, "lmcache_controller_manager", None
+        )
+
+        if controller_manager is None:
+            raise HTTPException(
+                status_code=503, detail="Controller manager not available"
+            )
+
+        reg_controller = controller_manager.reg_controller
+        registry = reg_controller.registry
+
+        # Get total key count
+        total_key_count = registry.get_total_kv_count()
+
+        # Get instances and their key counts
+        instances = []
+        total_instance_count = 0
+        total_worker_count = 0
+
+        for instance_id, instance_node in registry.instances.items():
+            total_instance_count += 1
+            workers = instance_node.workers.values()
+            num_workers = len(workers)
+            instance_key_count = sum(w.get_kv_count() for w in workers)
+            total_worker_count += num_workers
+            instances.append(
+                InstanceKeyStats(
+                    instance_id=instance_id,
+                    key_count=instance_key_count,
+                    worker_count=num_workers,
+                )
+            )
+
+        return KeyStatsResponse(
+            total_key_count=total_key_count,
+            total_instance_count=total_instance_count,
+            total_worker_count=total_worker_count,
+            instances=instances,
+        )
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e)) from None

lmcache/v1/internal_api_server/controller/worker_info_api.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Annotated, Optional
+
+# Third Party
+from fastapi import APIRouter, HTTPException, Query, Request
+from pydantic import BaseModel
+
+router = APIRouter()
+
+
+class WorkerInfoResponse(BaseModel):
+    instance_id: str
+    worker_id: int
+    ip: str
+    port: int
+    peer_init_url: Optional[str]
+    registration_time: float
+    last_heartbeat_time: float
+    key_count: int
+
+
+class WorkerListResponse(BaseModel):
+    workers: list[WorkerInfoResponse]
+    total_count: int
+
+
+@router.get("/controller/workers")
+async def get_workers(
+    request: Request,
+    instance_id: Annotated[Optional[str], Query()] = None,
+    worker_id: Annotated[Optional[int], Query()] = None,
+):
+    """
+    Get worker information with flexible query parameters.
+
+    - No parameters: List all registered workers across all instances
+    - instance_id only: List all workers for a specific instance
+    - instance_id and worker_id: Get detailed info about a specific worker
+
+    Args:
+        instance_id: Optional instance ID to filter workers
+        worker_id: Optional worker ID to get specific worker details
+    """
+    try:
+        controller_manager = getattr(
+            request.app.state, "lmcache_controller_manager", None
+        )
+
+        if controller_manager is None:
+            raise HTTPException(
+                status_code=503, detail="Controller manager not available"
+            )
+
+        reg_controller = controller_manager.reg_controller
+
+        # Case 1: Get specific worker by instance_id and worker_id
+        if instance_id is not None and worker_id is not None:
+            worker_node = reg_controller.registry.get_worker(instance_id, worker_id)
+            if worker_node is None:
+                raise HTTPException(
+                    status_code=404,
+                    detail=f"Worker ({instance_id}, {worker_id}) not found",
+                )
+
+            worker_info = worker_node.to_worker_info(instance_id)
+            key_count = worker_node.get_kv_count()
+            return WorkerInfoResponse(
+                instance_id=worker_info.instance_id,
+                worker_id=worker_info.worker_id,
+                ip=worker_info.ip,
+                port=worker_info.port,
+                peer_init_url=worker_info.peer_init_url,
+                registration_time=worker_info.registration_time,
+                last_heartbeat_time=worker_info.last_heartbeat_time,
+                key_count=key_count,
+            )
+
+        # Case 2: Get all workers for a specific instance
+        elif instance_id is not None:
+            instance_node = reg_controller.registry.get_instance(instance_id)
+            if instance_node is None:
+                raise HTTPException(
+                    status_code=404,
+                    detail=f"No workers found for instance {instance_id}",
+                )
+
+            worker_infos = instance_node.get_all_worker_infos()
+            workers = []
+            for worker_info in worker_infos:
+                worker_node = reg_controller.registry.get_worker(
+                    instance_id, worker_info.worker_id
+                )
+                key_count = worker_node.get_kv_count() if worker_node else 0
+                workers.append(
+                    WorkerInfoResponse(
+                        instance_id=worker_info.instance_id,
+                        worker_id=worker_info.worker_id,
+                        ip=worker_info.ip,
+                        port=worker_info.port,
+                        peer_init_url=worker_info.peer_init_url,
+                        registration_time=worker_info.registration_time,
+                        last_heartbeat_time=worker_info.last_heartbeat_time,
+                        key_count=key_count,
+                    )
+                )
+
+            return WorkerListResponse(workers=workers, total_count=len(workers))
+
+        # Case 3: Get all workers across all instances
+        else:
+            worker_infos = reg_controller.registry.get_all_worker_infos_cached()
+            workers = []
+            for worker_info in worker_infos:
+                worker_node = reg_controller.registry.get_worker(
+                    worker_info.instance_id, worker_info.worker_id
+                )
+                key_count = worker_node.get_kv_count() if worker_node else 0
+                workers.append(
+                    WorkerInfoResponse(
+                        instance_id=worker_info.instance_id,
+                        worker_id=worker_info.worker_id,
+                        ip=worker_info.ip,
+                        port=worker_info.port,
+                        peer_init_url=worker_info.peer_init_url,
+                        registration_time=worker_info.registration_time,
+                        last_heartbeat_time=worker_info.last_heartbeat_time,
+                        key_count=key_count,
+                    )
+                )
+
+            return WorkerListResponse(workers=workers, total_count=len(workers))
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e)) from None

lmcache/v1/internal_api_server/utils.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Dict, List
+
+
+def get_all_server_infos(config, worker_count) -> List[Dict[str, str]]:
+    """
+    Generate a list of server information (scheduler and workers) based on the config.
+
+    Args:
+        config: The configuration object containing server details.
+        worker_count: The number of worker servers.
+
+    Returns:
+        List[Dict[str, str]]: A JSON list with server information.
+    """
+    servers = []
+    include_index_list = getattr(config, "internal_api_server_include_index_list", None)
+    socket_path_prefix = getattr(config, "internal_api_server_socket_path_prefix", None)
+
+    # Add scheduler info (index 0)
+    if include_index_list is None or 0 in include_index_list:
+        port = config.internal_api_server_port_start
+        server_info = {
+            "name": f"{config.lmcache_instance_id}_scheduler",
+            "host": config.internal_api_server_host,
+            "port": f"{socket_path_prefix}_{port}" if socket_path_prefix else port,
+        }
+        servers.append(server_info)
+
+    # Add workers info (index 1 to worker_count)
+    for worker_id in range(worker_count):
+        port_offset = 1 + worker_id
+        if include_index_list is None or port_offset in include_index_list:
+            port = config.internal_api_server_port_start + port_offset
+            server_info = {
+                "name": f"{config.lmcache_instance_id}_worker{worker_id}",
+                "host": config.internal_api_server_host,
+                "port": f"{socket_path_prefix}_{port}" if socket_path_prefix else port,
+            }
+            servers.append(server_info)
+
+    return servers

lmcache/v1/internal_api_server/vllm/__init__.py

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