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

lmcache/v1/cache_controller/message.py

@@ -0,0 +1,828 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, List, Optional, Tuple, Union
+
+# Third Party
+import msgspec
+
+# First Party
+from lmcache.v1.cache_controller.commands.base import HeartbeatCommand
+from lmcache.v1.cache_controller.commands.full_sync import FullSyncCommand
+
+# Type alias for all command types - msgspec needs Union for tagged unions
+AnyCommand = Union[FullSyncCommand, HeartbeatCommand]
+
+
+class MsgBase(msgspec.Struct, tag=True):  # type: ignore
+    """Base class for all messages"""
+
+    def describe(self) -> str:
+        return ""
+
+
+# NOTE: The additional layer of abstraction is to
+# differentiate among
+# (1) WorkerMsg: push-pull (lmcache->controller)
+# (2) WorkerReqMsg: req-reply (lmcache->controller)
+# (3) ControlMessage: req-reply (controller->lmcache)
+# (4) OrchMsg: req-reply (ochestrator->controller)
+
+
+"""Message from LMCache to Controller"""
+
+
+class WorkerMsg(MsgBase):
+    """Message between LMCache and Controller"""
+
+    def describe(self) -> str:
+        return ""
+
+
+"""Worker Request (requiring a reply) Message from LMCache to Controller"""
+
+
+class WorkerReqMsg(MsgBase):
+    def describe(self) -> str:
+        return ""
+
+
+class RegisterMsg(WorkerReqMsg):
+    """Message for Registration (REQ-REP mode)"""
+
+    instance_id: str
+    worker_id: int
+    ip: str
+    port: int
+    # URL for actual KV cache transfer, only useful when p2p is enabled
+    peer_init_url: Optional[str]
+
+    def describe(self) -> str:
+        return (
+            f"Registering instance {self.instance_id}, "
+            f"worker {self.worker_id} "
+            f"at {self.ip}:{self.port}"
+            f" with peer init URL {self.peer_init_url}"
+        )
+
+
+class DeRegisterMsg(WorkerMsg):
+    """Message for Deregistration"""
+
+    instance_id: str
+    worker_id: int
+    ip: str
+    port: int
+
+    def describe(self) -> str:
+        return (
+            f"Deregistering instance {self.instance_id}, "
+            f"worker {self.worker_id} "
+            f"at {self.ip}:{self.port}"
+        )
+
+
+class KVOperationMsg(WorkerMsg):
+    """Base class for KV operation messages (admit/evict) with full context"""
+
+    instance_id: str
+    worker_id: int
+    key: int
+    location: str
+    seq_num: int = 0
+
+
+class KVAdmitMsg(KVOperationMsg):
+    """Message for KV chunk admission"""
+
+    def describe(self) -> str:
+        return f"kv_admit {self.key} to {self.instance_id}"
+
+
+class KVEvictMsg(KVOperationMsg):
+    """Message for KV chunk eviction"""
+
+    def describe(self) -> str:
+        return f"kv_evict {self.key} from {self.instance_id}"
+
+
+@dataclass
+class WorkerInfo:
+    """Information about a worker for external API consumption."""
+
+    instance_id: str
+    worker_id: int
+    ip: str
+    port: int
+    peer_init_url: Optional[str]
+    registration_time: float
+    last_heartbeat_time: float
+
+
+class OpType(Enum):
+    """Enum for KV operation types"""
+
+    ADMIT = "admit"
+    EVICT = "evict"
+
+
+class KVOpEvent(msgspec.Struct):
+    """Lightweight KV operation event for queue storage (without common fields)"""
+
+    op_type: OpType
+    key: int
+    seq_num: int
+
+
+class BatchedKVOperationMsg(WorkerMsg):
+    """Batched KV operation message with common fields and lightweight operations
+
+    Design: Common fields (instance_id, worker_id, location) are stored once
+    at the batch level, while individual operations only contain the varying
+    fields (op_type, key, seq_num). This reduces memory and network overhead.
+    """
+
+    instance_id: str
+    worker_id: int
+    location: str
+    operations: list[KVOpEvent]
+
+    def describe(self) -> str:
+        return (
+            f"Batched KV operations with {len(self.operations)} messages "
+            f"from {self.instance_id}:{self.worker_id}"
+        )
+
+
+# ============= Full Sync Messages (PUSH mode) =============
+
+
+class FullSyncBatchMsg(WorkerMsg):
+    """Full sync batch message (PUSH mode, no confirmation needed)
+
+    Used to send a batch of chunk hashes during full sync.
+    """
+
+    instance_id: str
+    worker_id: int
+    location: str
+    sync_id: str
+    batch_id: int  # Current batch number (0-indexed)
+    keys: list[int]  # List of chunk_hash in this batch
+
+    def describe(self) -> str:
+        return (
+            f"FullSyncBatch sync_id={self.sync_id} batch_id={self.batch_id} "
+            f"keys_count={len(self.keys)} from {self.instance_id}:{self.worker_id}"
+        )
+
+
+class FullSyncEndMsg(WorkerMsg):
+    """Full sync end message (PUSH mode)
+
+    Sent after all batches are sent to signal completion.
+    """
+
+    instance_id: str
+    worker_id: int
+    location: str
+    sync_id: str
+    actual_total_keys: int  # Actual total keys sent (for verification)
+
+    def describe(self) -> str:
+        return (
+            f"FullSyncEnd sync_id={self.sync_id} "
+            f"total_keys={self.actual_total_keys} from "
+            f"{self.instance_id}:{self.worker_id}"
+        )
+
+
+class HeartbeatMsg(WorkerReqMsg):
+    """Message for heartbeat (REQ-REP mode), include register info for re-register"""
+
+    instance_id: str
+    worker_id: int
+    ip: str
+    port: int
+    peer_init_url: Optional[str]
+
+    def describe(self) -> str:
+        return f"Heartbeat from instance {self.instance_id}, worker {self.worker_id}"
+
+
+class BatchedP2PLookupMsg(WorkerReqMsg):
+    """Batched P2P lookup message"""
+
+    hashes: list[int]
+    instance_id: str
+    worker_id: int  # TP rank
+
+    def describe(self) -> str:
+        return (
+            f"Batched P2P lookup for {len(self.hashes)} keys from "
+            f"instance id {self.instance_id} and "
+            f"worker id {self.worker_id}"
+        )
+
+
+class FullSyncStartMsg(WorkerReqMsg):
+    """Full sync start message (REQ-REP mode, needs confirmation)
+
+    Sent before starting full sync to notify controller and get confirmation.
+    """
+
+    instance_id: str
+    worker_id: int
+    location: str
+    sync_id: str  # Sync session ID
+    total_keys: int  # Expected total key count
+    batch_count: int  # Expected batch count
+
+    def describe(self) -> str:
+        return (
+            f"FullSyncStart sync_id={self.sync_id} "
+            f"total_keys={self.total_keys} batches={self.batch_count} "
+            f"from {self.instance_id}:{self.worker_id}"
+        )
+
+
+class FullSyncStatusMsg(WorkerReqMsg):
+    """Full sync status query message (REQ-REP mode)
+
+    Used to query sync progress and check if freeze mode can be exited.
+    """
+
+    instance_id: str
+    worker_id: int
+    sync_id: str
+
+    def describe(self) -> str:
+        return (
+            f"FullSyncStatus query sync_id={self.sync_id} "
+            f"from {self.instance_id}:{self.worker_id}"
+        )
+
+
+"""Worker Request Return Message from Controller back to LMCache"""
+
+
+class WorkerReqRetMsg(MsgBase):
+    def describe(self) -> str:
+        return ""
+
+
+class RegisterRetMsg(WorkerReqRetMsg):
+    """Register response message with controller configuration
+
+    Returns configuration information that the worker needs for initialization,
+    such as dedicated heartbeat URL.
+    """
+
+    # Extra configuration from controller to worker
+    # e.g., {"heartbeat_url": "tcp://...:8082"}
+    extra_config: dict[str, str] = msgspec.field(default_factory=dict)
+
+    def describe(self) -> str:
+        return f"RegisterRet extra_config={self.extra_config}"
+
+
+class HeartbeatRetMsg(WorkerReqRetMsg):
+    """Heartbeat response message with optional commands
+
+    The controller can use this to send commands to workers through
+    the heartbeat mechanism. This provides a general-purpose way to
+    trigger actions on workers without adding new message types.
+
+    The commands field uses msgspec's tagged union for polymorphic
+    serialization - each command subclass is identified by its tag.
+    Commands are executed sequentially by the worker.
+    """
+
+    commands: List[AnyCommand] = []
+
+    def describe(self) -> str:
+        if not self.commands:
+            return "HeartbeatRet (no commands)"
+        cmd_descs = [cmd.describe() for cmd in self.commands]
+        return f"HeartbeatRet commands=[{', '.join(cmd_descs)}]"
+
+    def has_commands(self) -> bool:
+        """Check if there are commands to execute"""
+        return len(self.commands) > 0
+
+
+class BatchedP2PLookupRetMsg(WorkerReqRetMsg):
+    """Batched P2P lookup return message"""
+
+    # (instance_id, location, num_hit_chunks, peer_init_url)
+    layout_info: list[tuple[str, str, int, str]]
+
+    def describe(self) -> str:
+        return f"The layout info is {self.layout_info}"
+
+
+class FullSyncStartRetMsg(WorkerReqRetMsg):
+    """Full sync start response message"""
+
+    sync_id: str
+    accepted: bool  # Whether sync request is accepted
+    error_msg: Optional[str] = None
+
+    def describe(self) -> str:
+        return f"FullSyncStartRet sync_id={self.sync_id} accepted={self.accepted}"
+
+
+class FullSyncStatusRetMsg(WorkerReqRetMsg):
+    """Full sync status response message"""
+
+    sync_id: str
+    is_complete: bool  # Whether current worker sync is complete
+    global_progress: float  # Global progress (0.0 - 1.0)
+    can_exit_freeze: bool  # Whether freeze mode can be exited
+    missing_batches: list[int] = []  # List of missing batch IDs that need to be resent
+
+    def describe(self) -> str:
+        missing_info = (
+            f" missing={self.missing_batches}" if self.missing_batches else ""
+        )
+        return (
+            f"FullSyncStatusRet sync_id={self.sync_id} "
+            f"complete={self.is_complete} progress={self.global_progress:.2%} "
+            f"can_exit_freeze={self.can_exit_freeze}{missing_info}"
+        )
+
+
+"""Control Message from Controller to LMCache"""
+
+
+class ControlMsg(MsgBase):
+    def describe(self) -> str:
+        return ""
+
+
+class ClearWorkerMsg(ControlMsg):
+    """Clear message for a single lmcache worker"""
+
+    worker_event_id: str
+    location: str
+
+    def describe(self) -> str:
+        return f"Clear tokens in location {self.location}"
+
+
+class PinWorkerMsg(ControlMsg):
+    """Pin message for a single lmcache worker"""
+
+    worker_event_id: str
+    location: str
+    tokens: list[int]
+
+    def describe(self) -> str:
+        return f"Pin tokens {self.tokens} in location {self.location}"
+
+
+class CompressWorkerMsg(ControlMsg):
+    """Compress message for a single lmcache worker"""
+
+    worker_event_id: str
+    method: str
+    location: str
+    tokens: Optional[list[int]] = None
+
+    def describe(self) -> str:
+        return (
+            f"Compress tokens {self.tokens} in "
+            f"locations {self.location} with "
+            f"method {self.method}"
+        )
+
+
+class DecompressWorkerMsg(ControlMsg):
+    """Decompress message for a single lmcache worker"""
+
+    worker_event_id: str
+    method: str
+    location: str
+    tokens: Optional[list[int]] = None
+
+    def describe(self) -> str:
+        return (
+            f"Decompress tokens {self.tokens} in "
+            f"locations {self.location} with "
+            f"method {self.method}"
+        )
+
+
+class MoveWorkerMsg(ControlMsg):
+    """Move message for a single lmcache worker"""
+
+    worker_event_id: str
+    old_position: str  # location (storage backend name)
+    new_position: Tuple[str, str]  # (target_url, location (storage backend name) )
+    tokens: Optional[list[int]] = None
+    copy: Optional[bool] = True
+
+    def describe(self) -> str:
+        return (
+            f"Move tokens {self.tokens} from {self.old_position} to {self.new_position}"
+        )
+
+
+class HealthWorkerMsg(ControlMsg):
+    """Health message for a single lmcache worker"""
+
+    worker_event_id: str
+
+    def describe(self) -> str:
+        return "Health check"
+
+
+class CheckFinishWorkerMsg(ControlMsg):
+    """Check finish message for a single lmcache worker"""
+
+    worker_event_id: str
+
+    def describe(self) -> str:
+        return f"Checking finish for worker event {self.worker_event_id}"
+
+
+class ControlRetMsg(MsgBase):
+    """Return message from LMCache to Controller"""
+
+    def describe(self) -> str:
+        return ""
+
+
+class ClearWorkerRetMsg(ControlRetMsg):
+    """Return message for a ClearWorkerMsg"""
+
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Number of cleared tokens: {self.num_tokens}"
+
+
+class PinWorkerRetMsg(ControlRetMsg):
+    """Pin return message for a single lmcache worker"""
+
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Number of pinned tokens: {self.num_tokens}"
+
+
+class CompressWorkerRetMsg(ControlRetMsg):
+    """Compress return message for a single lmcache worker"""
+
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Compress success: {self.num_tokens}"
+
+
+class DecompressWorkerRetMsg(ControlRetMsg):
+    """Decompress return message for a single lmcache worker"""
+
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Decompress success: {self.num_tokens}"
+
+
+class MoveWorkerRetMsg(ControlRetMsg):
+    """Move return message for a single lmcache worker"""
+
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Moving {self.num_tokens} tokens"
+
+
+class HealthWorkerRetMsg(ControlRetMsg):
+    """Health return message for a single lmcache worker"""
+
+    error_code: int
+
+    def describe(self) -> str:
+        return f"Health check error code: {self.error_code}"
+
+
+class CheckFinishWorkerRetMsg(ControlRetMsg):
+    """Check finish return message for a single lmcache worker"""
+
+    status: str
+
+    def describe(self) -> str:
+        return f"Check finish status: {self.status}"
+
+
+"""Orchestration Message from Ochestrator to LMCache"""
+
+
+class OrchMsg(MsgBase):
+    """Message from Ochestrator to Controller"""
+
+    def describe(self) -> str:
+        return ""
+
+
+class QueryInstMsg(OrchMsg):
+    """Query instance message"""
+
+    event_id: str
+    ip: str
+
+    def describe(self) -> str:
+        return f"Query instance id of ip {self.ip}"
+
+
+class LookupMsg(OrchMsg):
+    """Lookup message"""
+
+    event_id: str
+    tokens: list[int]
+
+    def describe(self) -> str:
+        return f"Lookup tokens {self.tokens}"
+
+
+class ClearMsg(OrchMsg):
+    """Clear message"""
+
+    event_id: str
+    instance_id: str
+    location: str
+
+    def describe(self) -> str:
+        return (
+            f"Clear tokens in instance {self.instance_id} and locations {self.location}"
+        )
+
+
+class PinMsg(OrchMsg):
+    """Pin message"""
+
+    event_id: str
+    instance_id: str
+    location: str
+    tokens: list[int]
+
+    def describe(self) -> str:
+        return (
+            f"Pin tokens {self.tokens} in instance "
+            f"{self.instance_id} and "
+            f"location {self.location}"
+        )
+
+
+class CompressMsg(OrchMsg):
+    """Compress message"""
+
+    event_id: str
+    instance_id: str
+    method: str
+    location: str
+    tokens: Optional[list[int]] = None  # `None` means compress all tokens
+
+    def describe(self) -> str:
+        return (
+            f"Compress tokens {self.tokens} in instance "
+            f"{self.instance_id} and "
+            f"locations {self.location} with "
+            f"method {self.method}"
+        )
+
+
+class DecompressMsg(OrchMsg):
+    """Decompress message"""
+
+    event_id: str
+    instance_id: str
+    method: str
+    location: str
+    tokens: Optional[list[int]] = None  # `None` means compress all tokens
+
+    def describe(self) -> str:
+        return (
+            f"Decompress tokens {self.tokens} in instance "
+            f"{self.instance_id} and "
+            f"locations {self.location} with "
+            f"method {self.method}"
+        )
+
+
+class MoveMsg(OrchMsg):
+    """Move message"""
+
+    event_id: str
+    old_position: Tuple[str, str]
+    new_position: Tuple[str, str]
+    tokens: Optional[list[int]] = None
+    copy: Optional[bool] = False
+
+    def describe(self) -> str:
+        return (
+            f"Move tokens {self.tokens} from {self.old_position} to {self.new_position}"
+        )
+
+
+class HealthMsg(OrchMsg):
+    """Health message"""
+
+    event_id: str
+    instance_id: str
+
+    def describe(self) -> str:
+        return f"Health check for instance {self.instance_id}"
+
+
+class CheckFinishMsg(OrchMsg):
+    """Check finish message"""
+
+    event_id: str
+
+    def describe(self) -> str:
+        return f"Checking finish for event {self.event_id}"
+
+
+class QueryWorkerInfoMsg(OrchMsg):
+    """Query worker info message"""
+
+    event_id: str
+    instance_id: str
+    worker_ids: Optional[list[int]]
+
+    def describe(self) -> str:
+        return f"Query worker info of {self.instance_id} : {self.worker_ids}"
+
+
+class OrchRetMsg(MsgBase):
+    """Return message from Controller to Ochestrator"""
+
+    def describe(self) -> str:
+        return ""
+
+
+class QueryInstRetMsg(OrchRetMsg):
+    """Query instance return message"""
+
+    event_id: str
+    instance_id: Optional[str]
+
+    def describe(self) -> str:
+        return f"The instance id is {self.instance_id}"
+
+
+class LookupRetMsg(OrchRetMsg):
+    """Lookup return message"""
+
+    event_id: str
+    layout_info: Dict[str, Tuple[str, int]]
+
+    def describe(self) -> str:
+        return f"The layout info is {self.layout_info}"
+
+
+class ClearRetMsg(OrchRetMsg):
+    """Clear return message"""
+
+    event_id: str
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Number of cleared tokens: {self.num_tokens}"
+
+
+class PinRetMsg(OrchRetMsg):
+    """Pin return message"""
+
+    event_id: str
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Number of pinned tokens: {self.num_tokens}"
+
+
+class CompressRetMsg(OrchRetMsg):
+    """Compress return message"""
+
+    event_id: str
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Compressed {self.num_tokens} tokens"
+
+
+class DecompressRetMsg(OrchRetMsg):
+    """Decompress return message"""
+
+    event_id: str
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Decompressed {self.num_tokens} tokens"
+
+
+class MoveRetMsg(OrchRetMsg):
+    """Move return message"""
+
+    event_id: str
+    num_tokens: int
+
+    def describe(self) -> str:
+        return f"Moving {self.num_tokens} tokens"
+
+
+class HealthRetMsg(OrchRetMsg):
+    """Health return message"""
+
+    event_id: str
+    # worker_id -> error_code
+    error_codes: Dict[int, int]
+
+    def describe(self) -> str:
+        return f"error_codes: {self.error_codes}"
+
+
+class CheckFinishRetMsg(OrchRetMsg):
+    """Check finish return message"""
+
+    status: str
+
+    def describe(self) -> str:
+        return f"Event status: {self.status}"
+
+
+class QueryWorkerInfoRetMsg(OrchRetMsg):
+    """Query worker info return message"""
+
+    event_id: str
+    worker_infos: list[WorkerInfo]
+
+    def describe(self) -> str:
+        return f"worker infos: {self.worker_infos}"
+
+
+class ErrorMsg(WorkerReqRetMsg):
+    """Control Error Message"""
+
+    error: str
+
+    def describe(self) -> str:
+        return f"Error: {self.error}"
+
+
+Msg = Union[
+    RegisterMsg,
+    RegisterRetMsg,
+    DeRegisterMsg,
+    KVAdmitMsg,
+    KVEvictMsg,
+    BatchedKVOperationMsg,
+    ClearWorkerMsg,
+    ClearWorkerRetMsg,
+    PinWorkerMsg,
+    PinWorkerRetMsg,
+    CompressWorkerMsg,
+    CompressWorkerRetMsg,
+    DecompressWorkerMsg,
+    DecompressWorkerRetMsg,
+    MoveWorkerMsg,
+    MoveWorkerRetMsg,
+    HealthWorkerMsg,
+    HealthWorkerRetMsg,
+    CheckFinishWorkerMsg,
+    CheckFinishWorkerRetMsg,
+    LookupMsg,
+    LookupRetMsg,
+    ClearMsg,
+    ClearRetMsg,
+    PinMsg,
+    PinRetMsg,
+    CompressMsg,
+    CompressRetMsg,
+    DecompressMsg,
+    DecompressRetMsg,
+    MoveMsg,
+    MoveRetMsg,
+    HealthMsg,
+    HealthRetMsg,
+    CheckFinishMsg,
+    CheckFinishRetMsg,
+    ErrorMsg,
+    QueryInstMsg,
+    QueryInstRetMsg,
+    HeartbeatMsg,
+    HeartbeatRetMsg,
+    BatchedP2PLookupMsg,
+    BatchedP2PLookupRetMsg,
+    QueryWorkerInfoMsg,
+    QueryWorkerInfoRetMsg,
+    FullSyncStartMsg,
+    FullSyncStartRetMsg,
+    FullSyncBatchMsg,
+    FullSyncEndMsg,
+    FullSyncStatusMsg,
+    FullSyncStatusRetMsg,
+]