wehub-resource-sync
110 lines
4.5 KiB
Python
110 lines
4.5 KiB
Python
# 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 IPCCacheServerKey, 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: IPCCacheServerKey - The key containing the token ids
|
|
# Returns: List of tuples (start, end) indicating the match ranges
|
|
"CB_LOOKUP_PRE_COMPUTED": ProtocolDefinition(
|
|
payload_classes=[IPCCacheServerKey],
|
|
response_class=list[tuple[int, int]],
|
|
handler_type=HandlerType.BLOCKING,
|
|
),
|
|
# Store pre-computed chunks
|
|
# Payload:
|
|
# - key: IPCCacheServerKey - 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=[IPCCacheServerKey, int, int, bytes],
|
|
response_class=tuple[bytes, bool],
|
|
handler_type=HandlerType.BLOCKING,
|
|
),
|
|
# Retrieve pre-computed chunks
|
|
# Payload:
|
|
# - key: IPCCacheServerKey - 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=[IPCCacheServerKey, list[tuple[int, int]], int, int, bytes],
|
|
response_class=tuple[bytes, bool],
|
|
handler_type=HandlerType.BLOCKING,
|
|
),
|
|
# Store final chunks after processing
|
|
# Payload:
|
|
# - key: IPCCacheServerKey - 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=[IPCCacheServerKey, 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,
|
|
),
|
|
}
|