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

lmcache/v1/multiprocess/protocol.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Main RPC protocol for the LMCache core server and clients.
+
+This module serves as the main entry point for the protocol system.
+All protocol definitions are now organized in the protocols/ subdirectory:
+- protocols/base.py: RequestType enum, HandlerType, ProtocolDefinition
+- protocols/engine.py: Core KV cache operations (REGISTER, STORE, RETRIEVE, etc.)
+- protocols/controller.py: Cache management operations (CLEAR, GET_CHUNK_SIZE)
+- protocols/debug.py: Debug and testing operations (NOOP)
+
+The protocol definitions are loaded and validated during initialization.
+"""
+
+# Standard
+from typing import Any, Optional
+
+# First Party
+from lmcache.v1.multiprocess.custom_types import IPCCacheEngineKey
+from lmcache.v1.multiprocess.protocols import initialize_protocols
+from lmcache.v1.multiprocess.protocols.base import HandlerType, RequestType
+
+# Initialize the protocol system
+# This loads all protocol definitions and validates them against the RequestType enum
+_PROTOCOL_DEFINITIONS = initialize_protocols()
+
+# Type aliases for backwards compatibility
+InstanceID = int
+KeyType = IPCCacheEngineKey
+
+
+def get_payload_classes(req_type: RequestType) -> list[Any]:
+    """
+    Get the expected payload classes for a request type.
+
+    Args:
+        req_type: The request type to look up
+
+    Returns:
+        List of expected payload classes in order
+
+    Raises:
+        ValueError: If the request type is not recognized
+    """
+    if pd := _PROTOCOL_DEFINITIONS.get(req_type, None):
+        return pd.payload_classes
+    else:
+        raise ValueError(f"Invalid request type: {req_type}")
+
+
+def get_response_class(req_type: RequestType) -> Optional[Any]:
+    """
+    Get the expected response class for a request type.
+
+    Args:
+        req_type: The request type to look up
+
+    Returns:
+        Expected response class, or None if no response
+
+    Raises:
+        ValueError: If the request type is not recognized
+    """
+    if pd := _PROTOCOL_DEFINITIONS.get(req_type, None):
+        return pd.response_class
+    else:
+        raise ValueError(f"Invalid request type: {req_type}")
+
+
+def get_handler_type(req_type: RequestType) -> HandlerType:
+    """
+    Get the handler type for a request type.
+
+    Args:
+        req_type: The request type to look up
+
+    Returns:
+        The handler type (SYNC, BLOCKING, or NON_BLOCKING)
+
+    Raises:
+        ValueError: If the request type is not recognized
+    """
+    if pd := _PROTOCOL_DEFINITIONS.get(req_type, None):
+        return pd.handler_type
+    else:
+        raise ValueError(f"Invalid request type: {req_type}")

lmcache/v1/multiprocess/protocols/README.md

@@ -0,0 +1,213 @@
+# Modular Protocol System
+
+This directory contains the modular protocol definitions for the LMCache multiprocess system.
+
+## Directory Structure
+
+```
+protocols/
+├── README.md           # This file
+├── __init__.py        # Protocol initialization and registration
+├── base.py            # Common types (HandlerType, ProtocolDefinition)
+├── engine.py          # Engine operations (STORE, RETRIEVE, etc.)
+├── controller.py      # Controller operations (CLEAR, GET_CHUNK_SIZE)
+└── debug.py           # Debug operations (NOOP)
+```
+
+## Design Overview
+
+The protocol system is designed to be modular, extensible, and IDE-friendly:
+
+1. **Static Enum with Validation**: The `RequestType` enum is defined statically in `base.py`:
+   - Provides perfect IDE autocomplete and type checking
+   - All request types visible to static analysis tools
+   - Validation ensures enum stays in sync with protocol definitions
+
+2. **Protocol Modules**: Each module (engine, controller, debug) defines:
+   - `REQUEST_NAMES`: List of request type names (for validation)
+   - `get_protocol_definitions()`: Returns dict of name → ProtocolDefinition
+
+3. **Validated Initialization**: The `__init__.py` module:
+   - Collects all protocol definitions from modules
+   - Validates each `RequestType` enum member has a definition
+   - Validates each definition has a corresponding enum member
+   - Ensures no duplicates or mismatches
+
+4. **Main Entry Point**: The `protocol.py` file:
+   - Calls `initialize_protocols()` on import
+   - Provides backwards-compatible API
+   - Exports `RequestType`, helper functions, and types
+
+## Adding New Protocols
+
+To add new protocol operations:
+
+### Option 1: Add to Existing Module
+
+If your operation fits an existing category (engine/controller/debug):
+
+1. **Add to the enum** in `protocols/base.py`:
+   ```python
+   class RequestType(enum.Enum):
+       # ... existing members ...
+       YOUR_NEW_OP = enum.auto()  # Add here
+   ```
+
+2. **Edit the appropriate protocol file** (e.g., `engine.py`):
+   - Add the request name to `REQUEST_NAMES`:
+     ```python
+     REQUEST_NAMES = [
+         "EXISTING_OP",
+         "YOUR_NEW_OP",  # Add here
+     ]
+     ```
+
+3. **Add the protocol definition** in `get_protocol_definitions()`:
+   ```python
+   def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+       return {
+           # ... existing definitions ...
+           "YOUR_NEW_OP": ProtocolDefinition(
+               payload_classes=[int, str],  # Your payload types
+               response_class=bool,          # Your response type
+               handler_type=HandlerType.SYNC,  # or BLOCKING
+           ),
+       }
+   ```
+
+4. **Done!** The validation system will verify everything matches on import.
+
+### Option 2: Create New Protocol Module
+
+If you're adding a new category of operations:
+
+1. **Add enum members** in `protocols/base.py`:
+   ```python
+   class RequestType(enum.Enum):
+       # ... existing members ...
+       
+       # Monitoring operations
+       HEALTH_CHECK = enum.auto()
+       GET_STATS = enum.auto()
+   ```
+
+2. **Create a new protocol file** (e.g., `monitoring.py`):
+   ```python
+   # SPDX-License-Identifier: Apache-2.0
+   """
+   Monitoring protocol definitions.
+   
+   This module defines protocols for:
+   - HEALTH_CHECK: Check server health status
+   - GET_STATS: Get cache statistics
+   """
+   from lmcache.v1.multiprocess.protocols.base import ProtocolDefinition, HandlerType
+
+   REQUEST_NAMES = [
+       "HEALTH_CHECK",
+       "GET_STATS",
+   ]
+
+   def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+       return {
+           "HEALTH_CHECK": ProtocolDefinition(
+               payload_classes=[],
+               response_class=dict,
+               handler_type=HandlerType.SYNC,
+           ),
+           "GET_STATS": ProtocolDefinition(
+               payload_classes=[],
+               response_class=dict,
+               handler_type=HandlerType.SYNC,
+           ),
+       }
+   ```
+
+3. **Register the module** in `__init__.py`:
+   ```python
+   from lmcache.v1.multiprocess.protocols import monitoring  # Import your module
+
+   def initialize_protocols():
+       protocol_modules = [
+           ("engine", engine),
+           ("controller", controller),
+           ("debug", debug),
+           ("monitoring", monitoring),  # Add here
+       ]
+       # ... rest of initialization ...
+   ```
+
+4. **Done!** The new operations are now available as:
+   - `RequestType.HEALTH_CHECK` (with full IDE autocomplete!)
+   - `RequestType.GET_STATS`
+
+## Using the Protocol System
+
+From any module in the codebase:
+
+```python
+from lmcache.v1.multiprocess.protocol import (
+    RequestType,
+    HandlerType,
+    get_payload_classes,
+    get_response_class,
+    get_handler_type,
+)
+
+# Use request types
+req_type = RequestType.STORE
+
+# Get protocol information
+payloads = get_payload_classes(req_type)
+response = get_response_class(req_type)
+handler = get_handler_type(req_type)
+```
+
+## Validation
+
+The initialization system validates at startup:
+
+1. **Enum-Definition Sync**: Every `RequestType` enum member must have a protocol definition
+2. **Definition-Enum Sync**: Every protocol definition must have a corresponding enum member
+3. **No duplicates**: Same request name cannot be defined in multiple modules
+4. **Complete definitions**: All names in `REQUEST_NAMES` must have definitions in `get_protocol_definitions()`
+
+If validation fails, `ProtocolInitializationError` is raised with a descriptive message pointing to the issue.
+
+### Example Error Messages
+
+```python
+# If you add RequestType.NEW_OP but forget the definition:
+ProtocolInitializationError: RequestType enum members {'NEW_OP'} have no protocol definitions.
+Add definitions to the appropriate protocol module or remove from the enum.
+
+# If you add a definition but forget the enum member:
+ProtocolInitializationError: Protocol definition 'NEW_OP' in module 'engine'
+has no corresponding RequestType enum member. Add 'RequestType.NEW_OP' to protocols/base.py
+```
+
+## Current Protocol Groups
+
+### Engine Operations (`engine.py`)
+Core KV cache operations:
+- `REGISTER_KV_CACHE`: Register a KV cache instance
+- `UNREGISTER_KV_CACHE`: Unregister a KV cache instance
+- `STORE`: Store KV cache blocks to the server
+- `RETRIEVE`: Retrieve KV cache blocks from the server
+- `LOOKUP`: Check if keys exist in the cache
+- `END_SESSION`: End a session and clean up resources
+
+### Controller Operations (`controller.py`)
+Cache management and configuration:
+- `CLEAR`: Clear all caches in the server
+- `GET_CHUNK_SIZE`: Get the chunk size configuration
+
+### Debug Operations (`debug.py`)
+Testing and monitoring:
+- `NOOP`: No-operation command for testing/heartbeat
+
+## Handler Types
+
+- `HandlerType.SYNC`: Fast operations that run directly in the main loop
+- `HandlerType.BLOCKING`: Operations that may block, run in a thread pool
+- `HandlerType.NON_BLOCKING`: Not yet supported (for future async handlers)

lmcache/v1/multiprocess/protocols/__init__.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Protocol initialization and registration system.
+
+This module provides the initialize_protocols() function that:
+1. Collects protocol definitions from all protocol modules
+2. Validates that the static RequestType enum matches protocol definitions
+3. Ensures all enum members have definitions and vice versa
+"""
+
+# First Party
+from lmcache.v1.multiprocess.protocols import (
+    blend,
+    blend_v2,
+    controller,
+    debug,
+    engine,
+    observability,
+)
+from lmcache.v1.multiprocess.protocols.base import (
+    HandlerType,
+    ProtocolDefinition,
+    RequestType,
+)
+
+
+class ProtocolInitializationError(Exception):
+    """Raised when there's an error during protocol initialization."""
+
+    pass
+
+
+_PROTOCOL_MODULES = [
+    ("engine", engine),
+    ("controller", controller),
+    ("debug", debug),
+    ("blend", blend),
+    ("blend_v2", blend_v2),
+    ("observability", observability),
+]
+
+
+def initialize_protocols() -> dict[RequestType, ProtocolDefinition]:
+    """
+    Initialize the protocol system by collecting all protocol definitions
+    and validating them against the RequestType enum.
+
+    This function:
+    1. Collects protocol definitions from all protocol modules
+    2. Validates that each RequestType enum member has a definition
+    3. Validates that each definition has a corresponding enum member
+    4. Ensures no duplicate or orphaned definitions
+
+    Returns:
+        protocol_definitions: Dict mapping RequestType enum values to
+        ProtocolDefinition
+
+    Raises:
+        ProtocolInitializationError: If there are mismatches between enum and
+        definitions
+    """
+    # Protocol modules to load
+    global _PROTOCOL_MODULES
+
+    # Step 1: Collect protocol definitions from all modules
+    protocol_definitions = {}
+    defined_names = set()
+    name_to_module: dict[str, str] = {}
+
+    for module_name, module in _PROTOCOL_MODULES:
+        module_defs = module.get_protocol_definitions()
+
+        # Check for duplicates across modules
+        for name in module_defs.keys():
+            if name in name_to_module:
+                raise ProtocolInitializationError(
+                    f"Duplicate protocol definition '{name}' found in modules "
+                    f"'{name_to_module[name]}' and '{module_name}'"
+                )
+            name_to_module[name] = module_name
+
+        # Validate that all names in REQUEST_NAMES have definitions
+        for name in module.REQUEST_NAMES:
+            if name not in module_defs:
+                raise ProtocolInitializationError(
+                    f"Request name '{name}' in module '{module_name}' "
+                    f"is listed in REQUEST_NAMES but has no protocol definition"
+                )
+            defined_names.add(name)
+
+        # Convert string names to enum values and store definitions
+        for name, definition in module_defs.items():
+            try:
+                enum_value = RequestType[name]
+                protocol_definitions[enum_value] = definition
+            except KeyError as err:
+                raise ProtocolInitializationError(
+                    f"Protocol definition '{name}' in module '{module_name}' "
+                    f"has no corresponding RequestType enum member. "
+                    f"Add 'RequestType.{name}' to protocols/base.py"
+                ) from err
+
+    # Step 2: Validate that all enum members have definitions
+    all_enum_names = {member.name for member in RequestType}
+    missing_definitions = all_enum_names - defined_names
+
+    if missing_definitions:
+        raise ProtocolInitializationError(
+            f"RequestType enum members {missing_definitions} have no protocol "
+            "definitions. Add definitions to the appropriate protocol module or "
+            "remove from the enum."
+        )
+
+    # Step 3: Validate that all definitions have enum members (already done in step 1)
+    # This is implicitly checked when we do RequestType[name]
+
+    return protocol_definitions
+
+
+# Export the base types for convenience
+__all__ = [
+    "initialize_protocols",
+    "RequestType",
+    "ProtocolDefinition",
+    "HandlerType",
+    "ProtocolInitializationError",
+]

lmcache/v1/multiprocess/protocols/base.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Base types and classes for the multiprocess protocol system.
+"""
+
+# Standard
+from dataclasses import dataclass
+from typing import Any, Optional
+import enum
+
+
+class HandlerType(enum.Enum):
+    """
+    Defines how a protocol handler should be executed.
+
+    - SYNC: Handler runs directly in the main loop (fast, non-blocking operations)
+    - BLOCKING: Handler may block, run in a thread pool (I/O, slow operations)
+    - NON_BLOCKING: Not supported yet (for future async handlers)
+    """
+
+    SYNC = enum.auto()
+    BLOCKING = enum.auto()
+    NON_BLOCKING = enum.auto()
+
+
+class RequestType(enum.Enum):
+    """
+    Enum of all available request types in the protocol system.
+
+    When adding a new request type:
+    1. Add the enum member here
+    2. Add the protocol definition in the appropriate protocols/*.py file
+    3. The validation system will ensure they stay in sync
+
+    Organized by category:
+    - Engine operations: Core KV cache operations
+    - Controller operations: Cache management and configuration
+    - Debug operations: Testing and monitoring
+    """
+
+    # Engine operations
+    REGISTER_KV_CACHE = enum.auto()
+    UNREGISTER_KV_CACHE = enum.auto()
+    STORE = enum.auto()
+    RETRIEVE = enum.auto()
+    LOOKUP = enum.auto()
+    QUERY_PREFETCH_STATUS = enum.auto()
+    QUERY_PREFETCH_LOOKUP_HITS = enum.auto()
+    FREE_LOOKUP_LOCKS = enum.auto()
+    END_SESSION = enum.auto()
+
+    # Controller operations
+    CLEAR = enum.auto()
+    GET_CHUNK_SIZE = enum.auto()
+    PING = enum.auto()
+
+    # Observability operations
+    REPORT_BLOCK_ALLOCATION = enum.auto()
+
+    # Debug operations
+    NOOP = enum.auto()
+
+    # Blend operations
+    CB_REGISTER_KV_CACHE = enum.auto()
+    CB_UNREGISTER_KV_CACHE = enum.auto()
+    CB_STORE_PRE_COMPUTED = enum.auto()
+    CB_LOOKUP_PRE_COMPUTED = enum.auto()
+    CB_RETRIEVE_PRE_COMPUTED = enum.auto()
+    CB_STORE_FINAL = enum.auto()
+
+    # Blend V2 operations (use CBMatchResult instead of list[tuple[int, int]])
+    CB_LOOKUP_PRE_COMPUTED_V2 = enum.auto()
+    CB_RETRIEVE_PRE_COMPUTED_V2 = enum.auto()
+
+
+@dataclass
+class ProtocolDefinition:
+    """
+    Defines the structure and behavior of a protocol request.
+
+    Attributes:
+        payload_classes: List of expected payload types in order
+        response_class: Expected response type, or None if no response
+        handler_type: How the handler should be executed (SYNC/BLOCKING/NON_BLOCKING)
+    """
+
+    payload_classes: list[Any]
+    response_class: Optional[Any]
+    handler_type: HandlerType

lmcache/v1/multiprocess/protocols/blend.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Debug protocol definitions for testing and monitoring.
+
+This module defines the protocol for:
+- NOOP: No-operation command for testing connectivity and as a heartbeat
+"""
+
+# First Party
+from lmcache.v1.multiprocess.custom_types import IPCCacheEngineKey, KVCache
+from lmcache.v1.multiprocess.protocols.base import HandlerType, ProtocolDefinition
+
+# Define request names for this protocol group
+REQUEST_NAMES = [
+    "CB_LOOKUP_PRE_COMPUTED",
+    "CB_STORE_PRE_COMPUTED",
+    "CB_RETRIEVE_PRE_COMPUTED",
+    "CB_STORE_FINAL",
+    "CB_REGISTER_KV_CACHE",
+    "CB_UNREGISTER_KV_CACHE",
+]
+
+
+def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+    """
+    Returns protocol definitions for debug operations.
+
+    Returns:
+        Dictionary mapping request names to their protocol definitions
+    """
+    return {
+        # Lookup pre-computed chunks
+        # Payload:
+        #   - key: IPCCacheEngineKey - The key containing the token ids
+        # Returns: List of tuples (start, end) indicating the match ranges
+        "CB_LOOKUP_PRE_COMPUTED": ProtocolDefinition(
+            payload_classes=[IPCCacheEngineKey],
+            response_class=list[tuple[int, int]],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Store pre-computed chunks
+        # Payload:
+        #   - key: IPCCacheEngineKey - The key containing the token ids
+        #   - offset: int - The starting offset in the CB KV cache buffer
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - event_ipc_handle: bytes - IPC handle for event notification
+        #                       when the pre-computed chunks are ready
+        # Returns:
+        #   - IPC handle bytes
+        #   - boolean flag indicating if the store is successful
+        "CB_STORE_PRE_COMPUTED": ProtocolDefinition(
+            payload_classes=[IPCCacheEngineKey, int, int, bytes],
+            response_class=tuple[bytes, bool],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Retrieve pre-computed chunks
+        # Payload:
+        #   - key: IPCCacheEngineKey - The key containing the token ids
+        #   - ranges: List[tuple[int, int]] - List of tuples (start, end) indicating
+        #                                     the match ranges to retrieve
+        #   - offset: int - The starting offset in the CB KV cache buffer
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - event_ipc_handle: bytes - IPC handle for event notification when the
+        #                       retrieval is complete
+        # Returns:
+        #   - IPC handle bytes
+        #   - boolean flag indicating if the retrieval is successful
+        "CB_RETRIEVE_PRE_COMPUTED": ProtocolDefinition(
+            payload_classes=[IPCCacheEngineKey, list[tuple[int, int]], int, int, bytes],
+            response_class=tuple[bytes, bool],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Store final chunks after processing
+        # Payload:
+        #   - key: IPCCacheEngineKey - The key containing the token ids
+        #   - offset: int - The starting offset in the CB KV cache buffer
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - event_ipc_handle: bytes - IPC handle for event notification
+        #                       when the final chunks are stored
+        # Returns:
+        #   - IPC handle bytes
+        #   - boolean flag indicating if the store is successful
+        "CB_STORE_FINAL": ProtocolDefinition(
+            payload_classes=[IPCCacheEngineKey, int, int, bytes],
+            response_class=tuple[bytes, bool],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Register CB KV Cache
+        # Payload:
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - kv_cache: KVCache - The CB KV cache configuration
+        #   - model_name: str - Name of the model associated with the engine
+        #   - world_size: int - World size of the engine
+        # Returns: None
+        "CB_REGISTER_KV_CACHE": ProtocolDefinition(
+            payload_classes=[int, KVCache, str, int],
+            response_class=None,
+            handler_type=HandlerType.SYNC,
+        ),
+        # Unregister CB KV Cache
+        # Payload:
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        # Returns: None
+        "CB_UNREGISTER_KV_CACHE": ProtocolDefinition(
+            payload_classes=[int],
+            response_class=None,
+            handler_type=HandlerType.SYNC,
+        ),
+    }

lmcache/v1/multiprocess/protocols/blend_v2.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Blend V2 protocol definitions.
+
+This module defines the V2 variants of blend lookup/retrieve protocols that use
+CBMatchResult instead of plain (start, end) integer tuples.  CBMatchResult carries
+old/cur position ranges plus the pre-computed chunk hash, enabling direct storage
+key lookup without re-hashing on the server side.
+"""
+
+# First Party
+from lmcache.v1.multiprocess.custom_types import CBMatchResult, IPCCacheEngineKey
+from lmcache.v1.multiprocess.protocols.base import HandlerType, ProtocolDefinition
+
+# Define request names for this protocol group
+REQUEST_NAMES = [
+    "CB_LOOKUP_PRE_COMPUTED_V2",
+    "CB_RETRIEVE_PRE_COMPUTED_V2",
+]
+
+
+def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+    """
+    Returns protocol definitions for blend V2 lookup/retrieve operations.
+
+    Returns:
+        Dictionary mapping request names to their protocol definitions
+    """
+    return {
+        # Lookup pre-computed chunks (V2)
+        # Payload:
+        #   - key: IPCCacheEngineKey - The key containing the token ids
+        # Returns: List of CBMatchResult with match positions and chunk hashes
+        "CB_LOOKUP_PRE_COMPUTED_V2": ProtocolDefinition(
+            payload_classes=[IPCCacheEngineKey],
+            response_class=list[CBMatchResult],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Retrieve pre-computed chunks (V2)
+        # Payload:
+        #   - key: IPCCacheEngineKey - The key containing the token ids
+        #   - cb_match_result: list[CBMatchResult] - Match results returned by
+        #                      CB_LOOKUP_PRE_COMPUTED_V2, with per-chunk hashes
+        #                      and query positions
+        #   - offset: int - The starting offset in the CB KV cache buffer
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - event_ipc_handle: bytes - IPC handle for event notification when the
+        #                       retrieval is complete
+        # Returns:
+        #   - IPC handle bytes
+        #   - boolean flag indicating if the retrieval is successful
+        "CB_RETRIEVE_PRE_COMPUTED_V2": ProtocolDefinition(
+            payload_classes=[IPCCacheEngineKey, list[CBMatchResult], int, int, bytes],
+            response_class=tuple[bytes, bool],
+            handler_type=HandlerType.BLOCKING,
+        ),
+    }