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

lmcache/v1/multiprocess/protocols/controller.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Controller protocol definitions for cache management and configuration.
+
+This module defines the protocol for:
+- CLEAR: Clear all caches in the server
+- GET_CHUNK_SIZE: Get the chunk size configuration from the server
+"""
+
+# First Party
+from lmcache.v1.multiprocess.protocols.base import HandlerType, ProtocolDefinition
+
+# Define request names for this protocol group
+REQUEST_NAMES = [
+    "CLEAR",
+    "GET_CHUNK_SIZE",
+    "PING",
+]
+
+
+def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+    """
+    Returns protocol definitions for controller operations.
+
+    Returns:
+        Dictionary mapping request names to their protocol definitions
+    """
+    return {
+        # Clear all caches
+        # Payload: None
+        # Returns: None
+        "CLEAR": ProtocolDefinition(
+            payload_classes=[],
+            response_class=None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Get chunk size configuration
+        # Payload: None
+        # Returns: int - The chunk size value
+        "GET_CHUNK_SIZE": ProtocolDefinition(
+            payload_classes=[],
+            response_class=int,
+            handler_type=HandlerType.SYNC,
+        ),
+        # Ping
+        # Payload: None
+        # Returns: bool - Always True
+        "PING": ProtocolDefinition(
+            payload_classes=[],
+            response_class=bool,
+            handler_type=HandlerType.BLOCKING,
+        ),
+    }

lmcache/v1/multiprocess/protocols/debug.py

@@ -0,0 +1,34 @@
+# 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.protocols.base import HandlerType, ProtocolDefinition
+
+# Define request names for this protocol group
+REQUEST_NAMES = [
+    "NOOP",
+]
+
+
+def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+    """
+    Returns protocol definitions for debug operations.
+
+    Returns:
+        Dictionary mapping request names to their protocol definitions
+    """
+    return {
+        # No-operation (for testing/heartbeat)
+        # Payload: None
+        # Returns: str - A confirmation message
+        "NOOP": ProtocolDefinition(
+            payload_classes=[],
+            response_class=str,
+            handler_type=HandlerType.SYNC,
+        ),
+    }

lmcache/v1/multiprocess/protocols/engine.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Engine protocol definitions for core KV cache operations.
+
+This module defines the protocol for:
+- REGISTER_KV_CACHE: Register a KV cache instance with the server
+- 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: Submit a prefix lookup and return a prefetch job ID
+- QUERY_PREFETCH_STATUS: Poll a prefetch job for its result
+- END_SESSION: End a session and clean up associated resources
+"""
+
+# First Party
+from lmcache.v1.gpu_connector.utils import LayoutHints
+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 = [
+    "REGISTER_KV_CACHE",
+    "UNREGISTER_KV_CACHE",
+    "STORE",
+    "RETRIEVE",
+    "LOOKUP",
+    "QUERY_PREFETCH_STATUS",
+    "QUERY_PREFETCH_LOOKUP_HITS",
+    "FREE_LOOKUP_LOCKS",
+    "END_SESSION",
+]
+
+# Type alias for cache keys
+KeyType = IPCCacheEngineKey
+
+
+def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+    """
+    Returns protocol definitions for engine operations.
+
+    Returns:
+        Dictionary mapping request names to their protocol definitions
+    """
+    return {
+        # Register KV Cache
+        # Payload:
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - kv_cache: KVCache - The KV cache configuration
+        #   - model_name: str - Name of the model associated with the engine
+        #   - world_size: int - World size of the engine
+        #   - layout_hints: LayoutHints - See custom_types.LayoutHints.
+        # Returns: None
+        "REGISTER_KV_CACHE": ProtocolDefinition(
+            payload_classes=[int, KVCache, str, int, LayoutHints],
+            response_class=None,
+            handler_type=HandlerType.SYNC,
+        ),
+        # Unregister KV Cache
+        # Payload:
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        # Returns: None
+        "UNREGISTER_KV_CACHE": ProtocolDefinition(
+            payload_classes=[int],
+            response_class=None,
+            handler_type=HandlerType.SYNC,
+        ),
+        # Store KV cache blocks
+        # Payload:
+        #   - key: KeyType - Cache key to store
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - gpu_block_ids: list[int] - GPU block IDs containing the data
+        #   - event_ipc_handle: bytes - CUDA event IPC handle for synchronization
+        # Returns: tuple[bytes, bool] - (CUDA event handle, success flag)
+        "STORE": ProtocolDefinition(
+            payload_classes=[KeyType, int, list[int], bytes],
+            response_class=tuple[bytes, bool],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Retrieve KV cache blocks
+        # Payload:
+        #   - key: KeyType - Cache key to retrieve
+        #   - instance_id: int - Unique identifier for the vLLM instance
+        #   - gpu_block_ids: list[int] - GPU block IDs to store retrieved data
+        #   - event_ipc_handle: bytes - CUDA event IPC handle for synchronization
+        #   - skip_first_n_tokens: int - Number of tokens to skip writing at the
+        #     start of the retrieve range (to avoid overwriting APC-shared blocks)
+        # Returns: tuple[bytes, bool] - (CUDA event handle, success flag)
+        "RETRIEVE": ProtocolDefinition(
+            payload_classes=[KeyType, int, list[int], bytes, int],
+            response_class=tuple[bytes, bool],
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Submit a prefix lookup; job is tracked server-side by request_id
+        # Payload:
+        #   - key: KeyType - Cache key to look up
+        #   - tp_size: int - Tensor-parallel size for
+        #       MLA multi-reader locking
+        # Returns: None
+        "LOOKUP": ProtocolDefinition(
+            payload_classes=[KeyType, int],
+            response_class=None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Query the status of a prefetch job by request_id
+        # Payload:
+        #   - request_id: str - The external request ID passed in the lookup key
+        # Returns: int | None - Chunk count when done, None if still in progress
+        "QUERY_PREFETCH_STATUS": ProtocolDefinition(
+            payload_classes=[str],
+            response_class=int | None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Query the lookup hit chunks before the prefetch is done
+        # Payload:
+        #   - request_id: str - The external request ID passed in the lookup key
+        # Returns: int | None - Chunk count if lookup is done, None if still in progress
+        "QUERY_PREFETCH_LOOKUP_HITS": ProtocolDefinition(
+            payload_classes=[str],
+            response_class=int | None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # Free locks (release read locks without a full RETRIEVE)
+        # Payload:
+        #   - key: KeyType - Cache key whose read locks
+        #       to release
+        #   - tp_size: int - Tensor-parallel size for
+        #       MLA multi-reader locking
+        # Returns: None
+        "FREE_LOOKUP_LOCKS": ProtocolDefinition(
+            payload_classes=[KeyType, int],
+            response_class=None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+        # End session
+        # Payload:
+        #   - request_id: str - Request ID of the session to end
+        # Returns: None
+        "END_SESSION": ProtocolDefinition(
+            payload_classes=[str],
+            response_class=None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+    }

lmcache/v1/multiprocess/protocols/observability.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Observability protocol definitions.
+
+This module defines protocols for:
+- REPORT_BLOCK_ALLOCATION: Report vLLM GPU block allocation events
+  (fire-and-forget, no response)
+"""
+
+# First Party
+from lmcache.v1.multiprocess.custom_types import BlockAllocationRecord
+from lmcache.v1.multiprocess.protocols.base import HandlerType, ProtocolDefinition
+
+# Define request names for this protocol group
+REQUEST_NAMES = [
+    "REPORT_BLOCK_ALLOCATION",
+]
+
+
+def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
+    """
+    Returns protocol definitions for observability operations.
+
+    Returns:
+        Dictionary mapping request names to their protocol definitions
+    """
+    return {
+        # Report vLLM block allocation
+        # Payload:
+        #   - instance_id: int - scheduler instance ID
+        #   - model_name: str - model name from the adapter
+        #   - records: list[BlockAllocationRecord] - allocation records
+        # Returns: None (fire-and-forget)
+        "REPORT_BLOCK_ALLOCATION": ProtocolDefinition(
+            payload_classes=[int, str, list[BlockAllocationRecord]],
+            response_class=None,
+            handler_type=HandlerType.BLOCKING,
+        ),
+    }