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

lmcache/v1/transfer_channel/py_socket_channel.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# TODO(baoloongmao): This module contains the control plane implementation
+# for socket-based channels. Currently, it provides the base class with
+# ZMQ-based handshake and initialization logic. In the future, this can
+# be extended to a complete channel with data plane using native sockets.
+
+# Standard
+from typing import Optional, Union
+import asyncio
+import threading
+import time
+
+# Third Party
+import msgspec
+import zmq
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.rpc_utils import get_zmq_context, get_zmq_socket
+from lmcache.v1.transfer_channel.abstract import BaseTransferChannel
+from lmcache.v1.transfer_channel.transfer_utils import (
+    InitSideMsgBase,
+    InitSideRetMsgBase,
+    SideMsg,
+)
+
+logger = init_logger(__name__)
+
+
+class PySocketMsgBase(msgspec.Struct, tag=True):
+    """Base class for all py-socket-related messages"""
+
+    pass
+
+
+class PySocketInitRequest(PySocketMsgBase):
+    """Initialization request, peer_init_url is used as peer identifier"""
+
+    peer_init_url: str
+
+
+class PySocketInitResponse(PySocketMsgBase):
+    """Initialization response"""
+
+    status: str
+
+
+PySocketMsg = Union[
+    PySocketInitRequest,
+    PySocketInitResponse,
+]
+
+
+class PySocketChannel(BaseTransferChannel):
+    """
+    Base class for socket-based transfer channels.
+
+    Control plane: Uses ZMQ sockets for handshake and initialization.
+    Data plane: To be implemented by subclasses.
+
+    Provides common control plane logic for different channel types.
+    """
+
+    def __init__(
+        self,
+        async_mode: bool = False,
+        **kwargs,
+    ):
+        assert "role" in kwargs
+        assert "buffer_ptr" in kwargs
+        assert "buffer_size" in kwargs
+        assert "align_bytes" in kwargs
+        assert "tp_rank" in kwargs
+        assert "peer_init_url" in kwargs
+
+        self.role = kwargs["role"]
+        self.buffer_ptr = kwargs["buffer_ptr"]
+        self.buffer_size = kwargs["buffer_size"]
+        self.align_bytes = kwargs["align_bytes"]
+        self.tp_rank = kwargs["tp_rank"]
+
+        self.peer_lookup_url = kwargs.get("peer_lookup_url", None)
+
+        self.running = True
+        self.remote_connections: dict[str, dict] = {}
+
+        self.side_channel: Optional[zmq.Socket] = None
+        self.running_threads: list[threading.Thread] = []
+
+        self.async_mode = async_mode
+        if self.async_mode:
+            self.zmq_context = get_zmq_context(use_asyncio=True)
+        else:
+            self.zmq_context = get_zmq_context(use_asyncio=False)
+        self.peer_init_url = kwargs["peer_init_url"]
+        self.event_loop = kwargs.get("event_loop", None)
+
+        self._init_side_channels()
+
+    ############################################################
+    # Control plane: Initialization functions
+    ############################################################
+
+    def lazy_init_peer_connection(
+        self,
+        local_id: str,
+        peer_id: str,
+        peer_init_url: str,
+        init_side_msg: Optional[InitSideMsgBase] = None,
+    ) -> Optional[InitSideRetMsgBase]:
+        raise NotImplementedError("Sync mode not supported in PySocketChannel")
+
+    async def async_lazy_init_peer_connection(
+        self,
+        local_id: str,
+        peer_id: str,
+        peer_init_url: str,
+        init_side_msg: Optional[InitSideMsgBase] = None,
+    ) -> Optional[InitSideRetMsgBase]:
+        """
+        Initialize connection to a peer using ZMQ sockets for handshake.
+
+        Note: peer_id is expected to be peer_init_url in this implementation.
+        """
+        init_tmp_socket = get_zmq_socket(
+            self.zmq_context,
+            peer_init_url,
+            "tcp",
+            zmq.REQ,
+            "connect",
+        )
+
+        init_req = PySocketInitRequest(peer_init_url=self.peer_init_url)
+        await init_tmp_socket.send(msgspec.msgpack.encode(init_req))
+
+        init_resp_bytes = await init_tmp_socket.recv()
+        _ = msgspec.msgpack.decode(init_resp_bytes, type=PySocketMsg)
+
+        self.remote_connections[peer_id] = {
+            "peer_init_url": peer_init_url,
+        }
+
+        init_ret_msg: Optional[InitSideRetMsgBase] = None
+        if init_side_msg is not None:
+            init_ret_msg = await self.async_send_init_side_msg(
+                init_tmp_socket,
+                init_side_msg,
+            )
+
+        init_tmp_socket.close()
+        return init_ret_msg
+
+    def _init_side_channels(self):
+        """Initialize side channel for handling incoming connections"""
+        if self.peer_init_url is None:
+            return
+
+        self.side_channel = get_zmq_socket(
+            self.zmq_context,
+            self.peer_init_url,
+            "tcp",
+            zmq.REP,
+            "bind",
+        )
+
+        if self.async_mode:
+            asyncio.run_coroutine_threadsafe(self._async_init_loop(), self.event_loop)
+        else:
+            self.init_thread = threading.Thread(target=self._init_loop, daemon=True)
+            self.init_thread.start()
+            self.running_threads.append(self.init_thread)
+
+    def _handle_init_msg(
+        self, req: Union[PySocketMsg, InitSideMsgBase]
+    ) -> Union[PySocketMsg, InitSideRetMsgBase]:
+        """Handle initialization messages from peers"""
+        resp: Union[PySocketMsg, InitSideRetMsgBase]
+        if isinstance(req, PySocketInitRequest):
+            peer_url = req.peer_init_url
+
+            self.remote_connections[peer_url] = {
+                "peer_init_url": peer_url,
+            }
+
+            self._on_peer_connected(peer_url)
+
+            resp = PySocketInitResponse(status="ok")
+            logger.info("Replying initialization response")
+
+        elif isinstance(req, InitSideMsgBase):
+            resp = self.handle_init_side_msg(req)
+            logger.info("Replying P2P init side response")
+        else:
+            raise ValueError(f"Unsupported InitMsg type: {type(req)}")
+
+        return resp
+
+    def _on_peer_connected(self, peer_url: str):
+        """Hook for subclasses to perform additional setup when a peer connects"""
+        pass
+
+    def _init_loop(self):
+        """Synchronous initialization loop for handling incoming connections"""
+        while self.running:
+            try:
+                req_bytes = self.side_channel.recv()
+
+                logger.info("Received initialization request")
+
+                req = msgspec.msgpack.decode(
+                    req_bytes, type=Union[PySocketMsg, SideMsg]
+                )
+
+                resp = self._handle_init_msg(req)
+
+                self.side_channel.send(msgspec.msgpack.encode(resp))
+
+            except Exception as e:
+                logger.error("Failed to process initialization loop: %s", str(e))
+                if self.running:
+                    time.sleep(0.01)
+
+    async def _async_init_loop(self):
+        """Asynchronous initialization loop for handling incoming connections"""
+        logger.info("Starting async initialization loop")
+
+        while self.running:
+            try:
+                req_bytes = await self.side_channel.recv()
+
+                logger.info("Received initialization request")
+
+                req = msgspec.msgpack.decode(
+                    req_bytes, type=Union[PySocketMsg, SideMsg]
+                )
+
+                resp = self._handle_init_msg(req)
+
+                await self.side_channel.send(msgspec.msgpack.encode(resp))
+
+            except Exception as e:
+                logger.error("Failed to process initialization loop: %s", str(e))
+                if self.running:
+                    await asyncio.sleep(0.01)
+
+    ############################################################
+    # Cleanup-related functions
+    ############################################################
+
+    def close(self):
+        """Close all sockets and cleanup resources"""
+        self.running = False
+        for thread in self.running_threads:
+            thread.join()
+
+        if self.side_channel is not None:
+            self.side_channel.close()
+
+        self.zmq_context.term()

lmcache/v1/transfer_channel/transfer_utils.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Union
+
+# Third Party
+import msgspec
+
+
+def get_correct_device(device: str, worker_id: int) -> str:
+    """
+    Get the correct device based on the given device string.
+
+    Args:
+        device (str): The device string, could be cpu or cuda.
+        worker_id (int): The worker id to determine the cuda device.
+
+    Returns:
+        str: The correct device string with device id.
+    """
+    if device == "cpu":
+        return "cpu"
+    elif device.startswith("cuda"):
+        return f"cuda:{worker_id}"
+    else:
+        raise ValueError(f"Invalid device: {device}")
+
+
+class SideMsgBase(msgspec.Struct, tag=True):
+    """Base class for all side-related messages during initialization"""
+
+    pass
+
+
+# Side messages during initialization
+class InitSideMsgBase(SideMsgBase):
+    """Base class for all side-related messages during initialization"""
+
+    pass
+
+
+class P2PInitSideMsg(InitSideMsgBase):
+    """P2P specific initialization message"""
+
+    pass
+
+
+# Side return messages during initialization
+class InitSideRetMsgBase(SideMsgBase):
+    """Base class for all side-related messages during initialization"""
+
+    pass
+
+
+class P2PInitSideRetMsg(InitSideRetMsgBase):
+    """P2P specific initialization return message"""
+
+    peer_lookup_url: str
+
+
+SideMsg = Union[
+    P2PInitSideMsg,
+    P2PInitSideRetMsg,
+]

lmcache/v1/utils/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

lmcache/v1/utils/bloom_filter.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import Union
+import array
+import hashlib
+import math
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class BloomFilter:
+    """Bloom Filter for memory-efficient."""
+
+    def __init__(
+        self, expected_elements: int = 1000000, false_positive_rate: float = 0.01
+    ):
+        self.size = self._optimal_size(expected_elements, false_positive_rate)
+        self.hash_count = self._optimal_hash_count(self.size, expected_elements)
+        array_size = (self.size + 31) // 32
+        self.bit_array = array.array("I", [0] * array_size)
+        self.item_count = 0
+        self.expected_elements = expected_elements
+        self.false_positive_rate = false_positive_rate
+
+    @staticmethod
+    def _optimal_size(n: int, p: float) -> int:
+        return int(-(n * math.log(p)) / (math.log(2) ** 2))
+
+    @staticmethod
+    def _optimal_hash_count(m: int, n: int) -> int:
+        return max(1, int((m / n) * math.log(2)))
+
+    def _hashes(self, item: Union[str, int]) -> list[int]:
+        result = []
+        if isinstance(item, int):
+            for i in range(self.hash_count):
+                h = hashlib.sha256(
+                    item.to_bytes((item.bit_length() + 7) // 8, "big", signed=False)
+                    + i.to_bytes(4, "big")
+                ).digest()
+                result.append(int.from_bytes(h[:4], "big") % self.size)
+        else:
+            for i in range(self.hash_count):
+                h = hashlib.sha256(f"{item}_{i}".encode()).digest()
+                result.append(int.from_bytes(h[:4], "big") % self.size)
+        return result
+
+    def add(self, item: Union[str, int]) -> None:
+        for pos in self._hashes(item):
+            self.bit_array[pos >> 5] |= 1 << (pos & 31)
+        self.item_count += 1
+
+    def add_batch_with_hashes_and_check(self, positions_list: list[list[int]]) -> int:
+        """Add multiple sets of hash positions and return count of new items.
+
+        Args:
+            positions_list: List of lists, where each inner list contains
+            hash positions for one item
+
+        Returns:
+            Number of new items that were actually added
+        """
+        unique_count = 0
+        for positions in positions_list:
+            is_new = any(
+                not (self.bit_array[pos >> 5] & (1 << (pos & 31))) for pos in positions
+            )
+            if is_new:
+                for pos in positions:
+                    self.bit_array[pos >> 5] |= 1 << (pos & 31)
+                unique_count += 1
+        self.item_count += unique_count
+        return unique_count
+
+    def contains(self, item: Union[str, int]) -> bool:
+        for pos in self._hashes(item):
+            if not (self.bit_array[pos >> 5] & (1 << (pos & 31))):
+                return False
+        return True
+
+    def clear(self) -> None:
+        self.bit_array = array.array("I", [0] * ((self.size + 31) // 32))
+        self.item_count = 0
+
+    def get_memory_usage_bytes(self) -> int:
+        return self.size // 8
+
+    def get_bit_set(self) -> int:
+        return sum(val.bit_count() for val in self.bit_array)
+
+    def get_fill_rate(self) -> float:
+        bits_set = self.get_bit_set()
+        return bits_set / self.size if self.size > 0 else 0.0
+
+    def get_statistics(self) -> dict:
+        bits_set = self.get_bit_set()
+        return {
+            "size_mb": self.get_memory_usage_bytes() / 1024 / 1024,
+            "hash_count": self.hash_count,
+            "item_count": self.item_count,
+            "bits_set": bits_set,
+            "fill_rate": bits_set / self.size if self.size > 0 else 0.0,
+            "expected_elements": self.expected_elements,
+            "false_positive_rate": self.false_positive_rate,
+        }

lmcache/v1/utils/cache_utils.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Callable, Generic, Optional, TypeVar
+import threading
+import time
+
+T = TypeVar("T")
+
+
+class TTLListCache(Generic[T]):
+    """
+    Generic list cache with TTL (Time-To-Live) support.
+
+    This cache stores a list of items of type T and automatically refreshes
+    when the cache expires.
+    """
+
+    def __init__(self, timeout_seconds: float = 30.0):
+        """
+        Initialize the cache.
+
+        Args:
+            timeout_seconds: Default TTL in seconds for cache entries.
+        """
+        self.cache: list[T] = []
+        self.cache_time: float = 0.0
+        self.cache_lock = threading.Lock()
+        self.timeout_seconds = timeout_seconds
+
+    def get_cached(
+        self,
+        get_fresh_data: Callable[[], list[T]],
+        timeout_override: Optional[float] = None,
+    ) -> list[T]:
+        """
+        Get cached data, refreshing if expired.
+
+        Args:
+            get_fresh_data: Function to get fresh data when cache is expired
+            timeout_override: Optional timeout override in seconds.
+                If None, use instance timeout_seconds.
+                If 0: cache immediately expires, always get fresh data and update cache.
+                If > 0: use specified timeout.
+
+        Returns:
+            List of cached items, either from cache or fresh data
+        """
+        current_time = time.time()
+
+        # Fast path: check cache without lock
+        if not self.is_expired(timeout_override, current_time):
+            return self.cache
+
+        # Slow path: refresh cache with lock
+        with self.cache_lock:
+            # Double-check after acquiring lock
+            if not self.is_expired(timeout_override, current_time):
+                return self.cache
+
+            # Refresh cache
+            self.cache = get_fresh_data()
+            self.cache_time = current_time
+            return self.cache
+
+    def clear(self) -> None:
+        """Clear the cache."""
+        with self.cache_lock:
+            self.cache.clear()
+            self.cache_time = 0.0
+
+    def is_expired(
+        self,
+        timeout_override: Optional[float] = None,
+        current_time: Optional[float] = None,
+    ) -> bool:
+        """
+        Check if the cache is expired without refreshing it.
+
+        Args:
+            timeout_override: Optional timeout override in seconds.
+                If None, use instance timeout_seconds.
+                If 0: always considered expired.
+                If > 0: use specified timeout.
+            current_time: Optional current time. If None, use time.time().
+
+        Returns:
+            True if cache is expired or uninitialized, False otherwise
+        """
+        if self.cache_time == 0.0:
+            return True
+
+        effective_timeout = (
+            timeout_override if timeout_override is not None else self.timeout_seconds
+        )
+        if current_time is None:
+            current_time = time.time()
+
+        return current_time - self.cache_time >= effective_timeout
+
+    def get_cache_age(self) -> float:
+        """
+        Get the age of the cache in seconds.
+
+        Returns:
+            Age in seconds. Returns float('inf') if cache is uninitialized.
+        """
+        if self.cache_time == 0.0:
+            return float("inf")
+        return time.time() - self.cache_time
+
+    def __len__(self) -> int:
+        """Get the number of items in cache (without checking expiration)."""
+        return len(self.cache)
+
+    def __repr__(self) -> str:
+        """String representation of the cache."""
+        age = self.get_cache_age()
+        if age == float("inf"):
+            age_str = "uninitialized"
+        else:
+            age_str = f"{age:.2f}s old"
+        return (
+            f"TTLListCache(cache_items={len(self)}, age={age_str}, "
+            f"timeout={self.timeout_seconds}s)"
+        )