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

lmcache/v1/storage_backend/batched_message_sender.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING, List, Optional
+import queue
+import threading
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.observability import PrometheusLogger
+from lmcache.v1.cache_controller.message import (
+    BatchedKVOperationMsg,
+    KVOpEvent,
+    OpType,
+)
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.metadata import LMCacheMetadata
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.cache_controller.worker import LMCacheWorker
+
+logger = init_logger(__name__)
+
+
+class BatchedMessageSender:
+    """
+    Batched message sender for KVOperation.
+
+    This class accumulates KV admit/evict messages and sends them in batches
+    to reduce communication overhead. Messages are flushed when either:
+    1. The batch size threshold is reached (default: 50 messages)
+    2. The timeout period expires (default: 0.01 seconds)
+
+    Each message is assigned a unique, monotonically increasing sequence number
+    to enable the receiver to detect missing or out-of-order messages.
+
+    Design rationale:
+    - Uses a SINGLE queue for both admit and evict messages to maintain strict
+      order consistency. This is critical because operations like
+      admit(key) -> evict(key) -> admit(key) must be processed in exact order
+      to avoid race conditions and state inconsistencies on the receiver side.
+
+    Thread-safe: Uses locks to protect internal queue and sequence counter.
+
+    Args:
+        metadata: Metadata for the worker
+        config: Configuration for the worker
+        location: Location of the worker
+        lmcache_worker: The worker to send messages to. If None, batching is disabled.
+    """
+
+    def __init__(
+        self,
+        metadata: LMCacheMetadata,
+        config: LMCacheEngineConfig,
+        location: str,
+        lmcache_worker: "LMCacheWorker",
+    ):
+        self.batch_size = config.get_extra_config_value("kv_msg_batch_size", 50)
+        self.batch_timeout = config.get_extra_config_value("kv_msg_batch_timeout", 0.01)
+        self.lmcache_worker = lmcache_worker
+
+        # Common fields shared by all operations in the batch
+        self.instance_id = config.lmcache_instance_id
+        self.worker_id = metadata.worker_id
+        self.location = location
+
+        # Use thread-safe queue for producer-consumer pattern
+        self.message_queue: queue.Queue[KVOpEvent] = queue.Queue()
+        self.sequence_number = 0
+        self.sequence_lock = threading.Lock()
+
+        # Condition variable for coordinating producer and consumer
+        self.cv = threading.Condition()
+        self.running = False
+        self.thread: Optional[threading.Thread] = None
+
+        self._start_background_thread()
+
+        self._setup_metrics()
+
+    def _setup_metrics(self):
+        """Setup metrics for monitoring queue size."""
+        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
+        if prometheus_logger is not None:
+            prometheus_logger.kv_msg_queue_size.set_function(
+                lambda: self.message_queue.qsize()
+            )
+
+    def _start_background_thread(self):
+        """Start background thread for periodic flushing."""
+        self.running = True
+        self.thread = threading.Thread(
+            target=self._consumer_loop, daemon=True, name="batched-msg-sender-thread"
+        )
+        self.thread.start()
+
+    def _consumer_loop(self):
+        """Consumer loop that drains queue and sends batched messages."""
+        while self.running:
+            with self.cv:
+                # Wait for timeout or notification from producer
+                self.cv.wait(timeout=self.batch_timeout)
+
+                # Check if we have messages to process while holding the lock
+                # This prevents race conditions but we'll release lock
+                # before blocking operations
+                if self.message_queue.empty():
+                    continue
+
+            # Drain the queue without holding the lock to avoid blocking producers
+            # This improves performance during the actual message processing
+            self._drain_and_send()
+
+    def _get_next_sequence_number(self) -> int:
+        """Get next sequence number for message tracking.
+
+        Thread-safe: Uses dedicated lock for sequence number generation.
+        """
+        with self.sequence_lock:
+            seq = self.sequence_number
+            self.sequence_number += 1
+            return seq
+
+    def add_kv_op(
+        self,
+        op_type: OpType,
+        key: int,
+    ):
+        """Add a KV operation to the batch queue.
+
+        Producer method: Adds operation to queue and notifies consumer
+        when batch size threshold is reached.
+
+        Args:
+            op_type: Operation type (ADMIT or EVICT)
+            key: Chunk hash key
+        """
+        # Create operation without sequence number (will be assigned during drain)
+        op = KVOpEvent(op_type=op_type, key=key, seq_num=-1)
+
+        # Thread-safe queue put
+        self.message_queue.put(op)
+
+        # Notify consumer if batch size threshold is reached
+        if self.message_queue.qsize() >= self.batch_size:
+            with self.cv:
+                self.cv.notify()
+
+    def _drain_and_send(self):
+        """Drain the queue and send all messages in a batch.
+
+        This method is called by the consumer thread to collect all pending
+        operations from the queue and send them as a single batched message.
+        """
+        ops_to_send: List[KVOpEvent] = []
+
+        # Drain all messages from the queue using blocking get with timeout
+        # This ensures we don't miss any messages due to race conditions
+        while True:
+            try:
+                # Use a small timeout to avoid blocking indefinitely
+                op = self.message_queue.get(timeout=0.001)
+                # Assign sequence number at drain time to ensure strict ordering
+                op.seq_num = self._get_next_sequence_number()
+                ops_to_send.append(op)
+            except queue.Empty:
+                # Queue is empty, break the loop
+                break
+
+        if not ops_to_send:
+            return
+
+        try:
+            # Ensure common fields are set
+            assert self.instance_id is not None, "instance_id must be set"
+            assert self.worker_id is not None, "worker_id must be set"
+            assert self.location is not None, "location must be set"
+
+            # Create batched message with common fields and lightweight operations
+            # This reduces redundancy: common fields are sent once instead of N times
+            batched_msg = BatchedKVOperationMsg(
+                instance_id=self.instance_id,
+                worker_id=self.worker_id,
+                location=self.location,
+                operations=ops_to_send,
+            )
+            self.lmcache_worker.put_msg(batched_msg)
+        finally:
+            # Mark all tasks as done regardless of success/failure
+            # This ensures flush() doesn't hang if put_msg fails
+            for _ in ops_to_send:
+                self.message_queue.task_done()
+
+    def flush(self):
+        """Manually flush all pending messages.
+
+        This method ensures all pending messages in the queue are processed
+        before returning. It triggers the consumer thread and waits for the
+        queue to be empty.
+        """
+        with self.cv:
+            self.cv.notify()
+
+        self.message_queue.join()
+
+    def close(self):
+        """Close the batched message sender and flush remaining messages."""
+        self.flush()
+        self.running = False
+
+        # Wake up consumer thread to exit
+        with self.cv:
+            self.cv.notify()
+
+        # Wait for thread to finish
+        if self.thread is not None and self.thread.is_alive():
+            self.thread.join(timeout=1.0)
+            if self.thread.is_alive():
+                logger.warning(
+                    "Batched message sender thread did not terminate within timeout"
+                )

lmcache/v1/storage_backend/cache_policy/__init__.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Dict, Type
+
+# First Party
+from lmcache.v1.storage_backend.cache_policy.base_policy import BaseCachePolicy
+from lmcache.v1.storage_backend.cache_policy.fifo import FIFOCachePolicy
+from lmcache.v1.storage_backend.cache_policy.lfu import LFUCachePolicy
+from lmcache.v1.storage_backend.cache_policy.lru import LRUCachePolicy
+from lmcache.v1.storage_backend.cache_policy.mru import MRUCachePolicy
+
+# Cache policy mapping
+POLICY_MAPPING: Dict[str, Type[BaseCachePolicy]] = {
+    "LRU": LRUCachePolicy,
+    "LFU": LFUCachePolicy,
+    "FIFO": FIFOCachePolicy,
+    "MRU": MRUCachePolicy,
+}
+
+
+def get_cache_policy(policy_name: str) -> BaseCachePolicy:
+    """
+    Factory function to get the cache policy instance based on the policy name.
+
+    Args:
+        policy_name: Name of the cache policy (case-insensitive, e.g., "LRU", "lru").
+
+    Returns:
+        Instance of the corresponding cache policy.
+
+    Raises:
+        ValueError: If the policy name is not supported.
+    """
+    if not policy_name:
+        raise ValueError("Cache policy name cannot be empty")
+
+    upper_policy_name = policy_name.upper()
+
+    try:
+        return POLICY_MAPPING[upper_policy_name]()
+    except KeyError:
+        raise ValueError(
+            f"Unknown cache policy: {upper_policy_name}."
+            f" Supported policies are: {list(POLICY_MAPPING.keys())}"
+        ) from None

lmcache/v1/storage_backend/cache_policy/base_policy.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from collections.abc import MutableMapping
+from typing import Generic, TypeVar
+import abc
+
+KeyType = TypeVar("KeyType")
+MapType = TypeVar("MapType", bound=MutableMapping)
+
+
+class BaseCachePolicy(Generic[KeyType, MapType], metaclass=abc.ABCMeta):
+    """
+    Interface for cache policy.
+    """
+
+    @abc.abstractmethod
+    def init_mutable_mapping(self) -> MapType:
+        """
+        Initialize a mutable mapping for cache storage.
+
+        Return:
+            A mutable mapping that can be used to store cache entries.
+        """
+        raise NotImplementedError
+
+    # TODO(Jiayi): we need to unify the `Any` type in the `MutableMapping`
+    @abc.abstractmethod
+    def update_on_hit(
+        self,
+        key: KeyType,
+        cache_dict: MapType,
+    ) -> None:
+        """
+        Update cache_dict and internal states when a cache is used
+
+        Input:
+            key: an object of KeyType
+            cache_dict: a dict consists of current cache
+        """
+        raise NotImplementedError
+
+    # TODO(Jiayi): we need to unify the `Any` type in the `MutableMapping`
+    @abc.abstractmethod
+    def update_on_put(
+        self,
+        key: KeyType,
+    ) -> None:
+        """
+        Update cache_dict and internal states when a cache is stored
+
+        Input:
+            key: an object of KeyType
+        """
+        raise NotImplementedError
+
+    # TODO(Jiayi): we need to unify the `Any` type in the `MutableMapping`
+    @abc.abstractmethod
+    def update_on_force_evict(
+        self,
+        key: KeyType,
+    ) -> None:
+        """
+        Update internal states when a cache is force evicted
+
+        Input:
+            key: an object of KeyType
+        """
+        raise NotImplementedError
+
+    # TODO(Jiayi): we need to unify the `Any` type in the `MutableMapping`
+    @abc.abstractmethod
+    def get_evict_candidates(
+        self,
+        cache_dict: MapType,
+        num_candidates: int = 1,
+    ) -> list[KeyType]:
+        """
+        Evict cache when a new cache comes and the storage is full
+
+        Input:
+            cache_dict: a dict consists of current cache
+            num_candidates: number of candidates to be evicted
+
+        Return:
+            return a list of keys to be evicted
+        """
+        raise NotImplementedError

lmcache/v1/storage_backend/cache_policy/fifo.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Any
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.storage_backend.cache_policy.base_policy import BaseCachePolicy, KeyType
+
+logger = init_logger(__name__)
+
+
+class FIFOCachePolicy(BaseCachePolicy[KeyType, dict[KeyType, Any]]):
+    """
+    FIFO cache policy.
+    """
+
+    def __init__(self):
+        logger.info("Initializing FIFOCachePolicy")
+
+    def init_mutable_mapping(self) -> dict[KeyType, Any]:
+        # NOTE(Jiayi): python dict maintains insertion order.
+        return {}
+
+    def update_on_hit(
+        self,
+        key: KeyType,
+        cache_dict: dict[KeyType, Any],
+    ) -> None:
+        pass
+
+    def update_on_put(
+        self,
+        key: KeyType,
+    ) -> None:
+        pass
+
+    def update_on_force_evict(
+        self,
+        key: KeyType,
+    ) -> None:
+        pass
+
+    # NOTE(Jiayi): We do best effort to get eviction candidates so the number
+    # of returned keys mignt be smaller than num_candidates.
+    def get_evict_candidates(
+        self,
+        cache_dict: dict[KeyType, Any],
+        num_candidates: int = 1,
+    ) -> list[KeyType]:
+        evict_keys = []
+        for key, cache in cache_dict.items():
+            if not cache.can_evict:
+                continue
+            evict_keys.append(key)
+            if len(evict_keys) == num_candidates:
+                break
+
+        return evict_keys

lmcache/v1/storage_backend/cache_policy/lfu.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import Any
+
+# Third Party
+from sortedcontainers import SortedDict
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.storage_backend.cache_policy.base_policy import BaseCachePolicy, KeyType
+
+logger = init_logger(__name__)
+
+
+class LFUCachePolicy(BaseCachePolicy[KeyType, dict[KeyType, Any]]):
+    """
+    LFU cache policy.
+    """
+
+    # NOTE(Jiayi): We use `sorted dict` + `bucket` to implement LFU.
+    # NOTE(Jiayi): We use FIFO for entries with the same frequency.
+    def __init__(self):
+        # TODO(Jiayi): `SortedDict` is log(N).
+        # A way to make it O(1) is to use a dict and keep track min freuency.
+        # However, this requires us keep another data structures to keep track
+        # of the pinned keys.
+        self.freq_to_keys = SortedDict()
+
+        # TODO(Jiayi): We can optimize this a bit by using `key_to_val_freq`
+        self.key_to_freq = {}
+
+        logger.info("Initializing LFUCachePolicy")
+
+    def init_mutable_mapping(self) -> dict[KeyType, Any]:
+        return {}
+
+    def update_on_hit(
+        self,
+        key: KeyType,
+        cache_dict: dict[KeyType, Any],
+    ) -> None:
+        curr_freq = self.key_to_freq[key]
+        self.freq_to_keys[curr_freq].pop(key)
+        if not self.freq_to_keys[curr_freq]:
+            self.freq_to_keys.pop(curr_freq)
+
+        curr_freq += 1
+        self.key_to_freq[key] = curr_freq
+
+        if curr_freq not in self.freq_to_keys:
+            self.freq_to_keys[curr_freq] = {key: None}
+        else:
+            self.freq_to_keys[curr_freq][key] = None
+
+    def update_on_put(
+        self,
+        key: KeyType,
+    ) -> None:
+        # Initialize the frequency for the new key.
+        self.key_to_freq[key] = 1
+        if 1 not in self.freq_to_keys:
+            self.freq_to_keys[1] = {key: None}
+        else:
+            self.freq_to_keys[1][key] = None
+
+    def update_on_force_evict(
+        self,
+        key: KeyType,
+    ) -> None:
+        freq = self.key_to_freq.pop(key, None)
+        if not freq:
+            return
+        self.freq_to_keys[freq].pop(key)
+        if not self.freq_to_keys[freq]:
+            self.freq_to_keys.pop(freq)
+
+    # NOTE(Jiayi): We do best effort to get eviction candidates so the number
+    # of returned keys mignt be smaller than num_candidates.
+    def get_evict_candidates(
+        self,
+        cache_dict: dict[KeyType, Any],
+        num_candidates: int = 1,
+    ) -> list[KeyType]:
+        evict_keys = []
+        evict_freqs = []
+        for curr_min_freq, fifo_keys in self.freq_to_keys.items():
+            for key in fifo_keys:
+                if not cache_dict[key].can_evict:
+                    continue
+                evict_keys.append(key)
+                evict_freqs.append(curr_min_freq)
+                self.key_to_freq.pop(key)
+                if len(evict_keys) == num_candidates:
+                    break
+
+            if len(evict_keys) == num_candidates:
+                break
+
+        for freq, key in zip(evict_freqs, evict_keys, strict=False):
+            self.freq_to_keys[freq].pop(key)
+            if not self.freq_to_keys[freq]:
+                self.freq_to_keys.pop(freq)
+
+        return evict_keys

lmcache/v1/storage_backend/cache_policy/lru.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from collections import OrderedDict
+from typing import Any, Dict
+import time
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.observability import LMCStatsMonitor
+from lmcache.utils import CacheEngineKey
+from lmcache.v1.storage_backend.cache_policy.base_policy import BaseCachePolicy, KeyType
+
+logger = init_logger(__name__)
+
+
+class LRUCachePolicy(BaseCachePolicy[KeyType, OrderedDict[KeyType, Any]]):
+    """
+    LRU cache policy.
+    """
+
+    def __init__(self):
+        logger.info("Initializing LRUCachePolicy")
+        self.chunk_hash_to_init_timestamp: Dict[Any, float] = {}
+        self.stats_monitor = LMCStatsMonitor.GetOrCreate()
+        self.max_num_chunk_hash = 12500000
+
+    def init_mutable_mapping(self) -> OrderedDict[KeyType, Any]:
+        return OrderedDict()
+
+    def update_chunk_hash_dict(self, key: KeyType) -> None:
+        curr_time = time.time()
+        # HACK: doing type conversion here
+        key_hash: Any = key
+        if isinstance(key, CacheEngineKey):
+            key_hash = key.chunk_hash
+
+        if init_timestamp := self.chunk_hash_to_init_timestamp.get(key_hash, None):
+            time_interval = curr_time - init_timestamp
+            self.stats_monitor.on_chunk_reuse(time_interval)
+        else:
+            if len(self.chunk_hash_to_init_timestamp) >= self.max_num_chunk_hash:
+                self.chunk_hash_to_init_timestamp.clear()
+            self.chunk_hash_to_init_timestamp[key_hash] = curr_time
+
+    def update_on_hit(
+        self,
+        key: KeyType,
+        cache_dict: OrderedDict[KeyType, Any],
+    ) -> None:
+        self.update_chunk_hash_dict(key)
+        cache_dict.move_to_end(key)
+
+    def update_on_put(
+        self,
+        key: KeyType,
+    ) -> None:
+        self.update_chunk_hash_dict(key)
+        pass
+
+    def update_on_force_evict(
+        self,
+        key: KeyType,
+    ) -> None:
+        pass
+
+    # NOTE(Jiayi): We do best effort to get eviction candidates so the number
+    # of returned keys mignt be smaller than num_candidates.
+    def get_evict_candidates(
+        self,
+        cache_dict: OrderedDict[KeyType, Any],
+        num_candidates: int = 1,
+    ) -> list[KeyType]:
+        evict_keys = []
+        for key, cache in cache_dict.items():
+            if not cache.can_evict:
+                continue
+            evict_keys.append(key)
+            if len(evict_keys) == num_candidates:
+                break
+
+        return evict_keys

lmcache/v1/storage_backend/cache_policy/mru.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from collections import OrderedDict
+from typing import Any
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.storage_backend.cache_policy.base_policy import BaseCachePolicy, KeyType
+
+logger = init_logger(__name__)
+
+
+class MRUCachePolicy(BaseCachePolicy[KeyType, OrderedDict[KeyType, Any]]):
+    """
+    MRU cache policy.
+    """
+
+    def __init__(self):
+        logger.info("Initializing MRUCachePolicy")
+
+    def init_mutable_mapping(self) -> OrderedDict[KeyType, Any]:
+        return OrderedDict()
+
+    def update_on_hit(
+        self,
+        key: KeyType,
+        cache_dict: OrderedDict[KeyType, Any],
+    ) -> None:
+        # since MRU evicts from the back, the logic is same as LRU.
+        cache_dict.move_to_end(key, last=True)
+
+    def update_on_put(
+        self,
+        key: KeyType,
+    ) -> None:
+        # No action needed for MRU on put, as the key is already at the back.
+        pass
+
+    def update_on_force_evict(
+        self,
+        key: KeyType,
+    ) -> None:
+        pass
+
+    # NOTE(Jiayi): We do best effort to get eviction candidates so the number
+    # of returned keys mignt be smaller than num_candidates.
+    def get_evict_candidates(
+        self,
+        cache_dict: OrderedDict[KeyType, Any],
+        num_candidates: int = 1,
+    ) -> list[KeyType]:
+        evict_keys = []
+        # Since the most recent object is at the end, we reverse the order here
+        for key, cache in reversed(cache_dict.items()):
+            if not cache.can_evict:
+                continue
+            evict_keys.append(key)
+            if len(evict_keys) == num_candidates:
+                break
+
+        return evict_keys