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

lmcache/v1/rpc_utils.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Literal, Optional
+import os
+import socket
+
+# Third Party
+import zmq
+import zmq.asyncio
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+ServiceKind = Literal["lookup", "offload", "lookup_worker", "lookup_scheduler"]
+
+# Default timeout constants for socket operations (in milliseconds)
+DEFAULT_SOCKET_RECV_TIMEOUT_MS = 30000
+DEFAULT_SOCKET_SEND_TIMEOUT_MS = 10000
+
+
+def get_zmq_context(use_asyncio: bool = True):
+    if use_asyncio:
+        return zmq.asyncio.Context.instance()
+    else:
+        return zmq.Context.instance()
+
+
+def get_zmq_socket(
+    context, socket_path: str, protocol: str, role, bind_or_connect: str
+):
+    """
+    Create a ZeroMQ socket with the specified protocol and role.
+    """
+    socket_addr = f"{protocol}://{socket_path}"
+    socket = context.socket(role)
+    if bind_or_connect == "bind":
+        socket.bind(socket_addr)
+    elif bind_or_connect == "connect":
+        socket.connect(socket_addr)
+    else:
+        raise ValueError(f"Invalid bind_or_connect: {bind_or_connect}")
+
+    return socket
+
+
+def get_zmq_socket_with_timeout(
+    context,
+    socket_path: str,
+    protocol: str,
+    role,
+    bind_or_connect: str,
+    recv_timeout_ms: int,
+    send_timeout_ms: int,
+) -> zmq.asyncio.Socket:
+    """
+    Create a ZeroMQ socket with timeout settings.
+    """
+    socket = get_zmq_socket(
+        context,
+        socket_path,
+        protocol,
+        role,
+        bind_or_connect,
+    )
+    # Only set RCVTIMEO for client role connect sockets
+    if bind_or_connect == "connect":
+        socket.setsockopt(zmq.RCVTIMEO, recv_timeout_ms)
+    socket.setsockopt(zmq.SNDTIMEO, send_timeout_ms)
+    return socket
+
+
+def close_zmq_socket(socket: zmq.asyncio.Socket, linger: int = 0) -> None:
+    """
+    Close a ZeroMQ socket cleanly.
+
+    :param socket: The zmq.Socket to be closed.
+    :param linger: LINGER period (in milliseconds).
+    Default is 0 (drop immediately).
+    """
+    try:
+        socket.setsockopt(zmq.LINGER, linger)  # type: ignore[attr-defined]
+        socket.close()
+    except Exception as e:
+        logger.error(f"Warning: Failed to close socket cleanly: {e}")
+
+
+def get_ip():
+    """
+    Get the local IP address of the machine.
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        # "Connect" to a public IP — just to determine local IP
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except Exception:
+        logger.warning(
+            "Failed to get local IP address. Falling back to loopback address."
+        )
+        return "127.0.0.1"  # Fallback to loopback
+    finally:
+        s.close()
+
+
+def get_zmq_rpc_path_lmcache(
+    engine_id: str,
+    service_name: ServiceKind = "lookup",
+    rpc_port: int = 0,
+    rank: int = 0,
+    base_url: Optional[str] = None,
+) -> str:
+    """Get the ZMQ RPC path for LMCache lookup and offload communication.
+
+    Args:
+        engine_id: The engine ID for the RPC path.
+        service_name: The service name, one of 'lookup', 'offload',
+            'lookup_worker', 'lookup_scheduler'.
+        rpc_port: The RPC port number.
+        rank: The rank of the worker.
+        base_url: Optional base URL for the socket path. If not provided,
+            will try to use VLLM_RPC_BASE_PATH or fallback to /tmp/vllm_rpc.
+
+    Returns:
+        The ZMQ socket path string.
+    """
+    if base_url is None:
+        # Try to import vllm.envs, fallback to default if not available
+        try:
+            # Third Party
+            import vllm.envs as envs
+
+            base_url = envs.VLLM_RPC_BASE_PATH
+        except (ImportError, ModuleNotFoundError):
+            # Fallback for testing environments without vllm
+            base_url = "/tmp/vllm_rpc"
+            logger.debug("vllm not available, using default base_url: %s", base_url)
+            # Ensure the directory exists for IPC socket
+            os.makedirs(base_url, exist_ok=True)
+
+    if service_name not in {"lookup", "offload", "lookup_worker", "lookup_scheduler"}:
+        raise ValueError(
+            f"service_name must be 'lookup' or 'offload', got {service_name!r}"
+        )
+
+    if isinstance(rpc_port, str):
+        rpc_port = rpc_port + str(rank)
+    else:
+        rpc_port += rank
+
+    logger.debug(
+        "Base URL: %s, Engine: %s, Service Name: %s, RPC Port: %s",
+        base_url,
+        engine_id,
+        service_name,
+        rpc_port,
+    )
+
+    socket_path = (
+        f"{base_url}/engine_{engine_id}_service_{service_name}_"
+        f"lmcache_rpc_port_{rpc_port}"
+    )
+
+    return socket_path

lmcache/v1/server/__init__.py

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

lmcache/v1/server/__main__.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import socket
+import threading
+import time
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.memory_management import MemoryFormat
+from lmcache.v1.protocol import (
+    ClientCommand,
+    ClientMetaMessage,
+    ServerMetaMessage,
+    ServerReturnCode,
+)
+from lmcache.v1.server.storage_backend import CreateStorageBackend
+
+logger = init_logger(__name__)
+
+
+class LMCacheServer:
+    def __init__(self, host, port, device):
+        self.host = host
+        self.port = port
+        # self.data_store = {}
+        self.data_store = CreateStorageBackend(device)
+        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.server_socket.bind((host, port))
+        self.server_socket.listen()
+
+    def receive_all(self, client_socket, n):
+        data = bytearray()
+        while len(data) < n:
+            packet = client_socket.recv(n - len(data))
+            if not packet:
+                return None
+            data.extend(packet)
+        return data
+
+    def handle_client(self, client_socket):
+        try:
+            while True:
+                logger.debug("Waiting for command")
+                header = self.receive_all(client_socket, ClientMetaMessage.packlength())
+                if not header:
+                    break
+                meta = ClientMetaMessage.deserialize(header)
+                logger.debug(f"Received command: {meta.command}")
+                match meta.command:
+                    case ClientCommand.PUT:
+                        t0 = time.perf_counter()
+                        s = self.receive_all(client_socket, meta.length)
+                        t1 = time.perf_counter()
+                        self.data_store.put(meta, s)
+                        t2 = time.perf_counter()
+                        logger.debug(
+                            f"Time to receive data: {t1 - t0}, time to store "
+                            f"data: {t2 - t1}"
+                        )
+
+                    case ClientCommand.GET:
+                        t0 = time.perf_counter()
+                        lms_memory_obj = self.data_store.get(meta.key)
+                        t1 = time.perf_counter()
+                        if lms_memory_obj is not None:
+                            client_socket.sendall(
+                                ServerMetaMessage(
+                                    ServerReturnCode.SUCCESS,
+                                    lms_memory_obj.length,
+                                    lms_memory_obj.fmt,
+                                    lms_memory_obj.dtype,
+                                    lms_memory_obj.shape,
+                                ).serialize()
+                            )
+                            t2 = time.perf_counter()
+                            client_socket.sendall(lms_memory_obj.data)
+                            t3 = time.perf_counter()
+                            logger.debug(
+                                f"Time to get data: {t1 - t0}, time to send "
+                                f"meta: {t2 - t1}, time to send data: {t3 - t2}"
+                            )
+                        else:
+                            client_socket.sendall(
+                                ServerMetaMessage(
+                                    ServerReturnCode.FAIL,
+                                    0,
+                                    MemoryFormat(1),
+                                    torch.float16,
+                                    torch.Size((0, 0, 0, 0)),
+                                ).serialize()
+                            )
+
+                    case ClientCommand.EXIST:
+                        code = (
+                            ServerReturnCode.SUCCESS
+                            if self.data_store.contains(meta.key)
+                            else ServerReturnCode.FAIL
+                        )
+                        logger.debug(f"Key exists: {code}")
+                        client_socket.sendall(
+                            ServerMetaMessage(
+                                code,
+                                0,
+                                MemoryFormat(1),
+                                torch.float16,
+                                torch.Size((0, 0, 0, 0)),
+                            ).serialize()
+                        )
+                    case ClientCommand.HEALTH:
+                        client_socket.sendall(
+                            ServerMetaMessage(
+                                ServerReturnCode.SUCCESS,
+                                0,
+                                MemoryFormat(1),
+                                torch.float16,
+                                torch.Size((0, 0, 0, 0)),
+                            ).serialize()
+                        )
+                        logger.debug("Health check successful")
+
+                    # TODO(Jiayi): Implement List
+                    # case ClientCommand.LIST:
+                    #     keys = list(self.data_store.list_keys())
+                    #     data = "\n".join(keys).encode()
+                    #     client_socket.sendall(
+                    #         ServerMetaMessage(ServerReturnCode.SUCCESS,
+                    #                           len(data)).serialize())
+                    #     client_socket.sendall(data)
+
+        finally:
+            logger.info("Client disconnected")
+            client_socket.close()
+
+    def run(self):
+        logger.info(f"Server started at {self.host}:{self.port}")
+        try:
+            while True:
+                client_socket, addr = self.server_socket.accept()
+                logger.info(f"Connected by {addr}")
+                threading.Thread(
+                    target=self.handle_client, args=(client_socket,)
+                ).start()
+        finally:
+            self.server_socket.close()
+
+
+def main():
+    # Standard
+    import sys
+
+    if len(sys.argv) not in [3, 4]:
+        logger.error(f"Usage: {sys.argv[0]} <host> <port> <storage>(default:cpu)")
+        exit(1)
+
+    host = sys.argv[1]
+    port = int(sys.argv[2])
+    if len(sys.argv) == 4:
+        device = sys.argv[3]
+    else:
+        device = "cpu"
+
+    server = LMCacheServer(host, port, device)
+    server.run()
+
+
+if __name__ == "__main__":
+    main()

lmcache/v1/server/storage_backend/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.server.storage_backend.abstract_backend import LMSBackendInterface
+from lmcache.v1.server.storage_backend.local_backend import LMSLocalBackend
+
+logger = init_logger(__name__)
+
+
+def CreateStorageBackend(device: str) -> LMSBackendInterface:
+    match device:
+        case "cpu":
+            # cpu only
+            logger.info("Initializing cpu-only cache server")
+            return LMSLocalBackend()
+        case _:
+            raise ValueError(f"Unsupported device: {device}")
+        # TODO(Jiayi): please implement hierarchical remote storage
+        # case _:
+        #    logger.info("Initializing disk-only cache server")
+        #    return LMSLocalDiskBackend(path=device)

lmcache/v1/server/storage_backend/abstract_backend.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import List, Optional
+import abc
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.utils import CacheEngineKey
+from lmcache.v1.protocol import ClientMetaMessage
+from lmcache.v1.server.utils import LMSMemoryObj
+
+logger = init_logger(__name__)
+
+
+class LMSBackendInterface(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def put(
+        self,
+        client_meta: ClientMetaMessage,
+        kv_chunk_bytes: bytearray,
+    ) -> None:
+        """
+        Store the KV cache of the tokens into the cache server.
+
+        Args:
+            key: the key of the token chunk, in the format of CacheEngineKey
+            client_meta: metadata sent by the client
+            kv_chunk_bytes: the kv cache (bytearray) of the token chunk
+
+        Returns:
+            None
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def contains(
+        self,
+        key: CacheEngineKey,
+    ) -> bool:
+        """
+        Query if a key is in the cache or not
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get(
+        self,
+        key: CacheEngineKey,
+    ) -> Optional[LMSMemoryObj]:
+        """
+        Retrieve LMSMemoryObj by the given key
+
+        Input:
+            key: the CacheEngineKey
+
+        Output:
+            An LMSMemoryObj object that contains the KV cache bytearray
+            with the some metadata
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def list_keys(
+        self,
+    ) -> List[CacheEngineKey]:
+        """
+        List all keys in the cache server
+
+        Output:
+            All keys in the cache server
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close(self):
+        """
+        Do the cleanup things
+        Children classes should override this method if necessary
+        """
+        pass

lmcache/v1/server/storage_backend/local_backend.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from collections import OrderedDict
+from typing import List, Optional
+import threading
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.utils import CacheEngineKey, _lmcache_nvtx_annotate
+from lmcache.v1.protocol import ClientMetaMessage
+from lmcache.v1.server.storage_backend.abstract_backend import LMSBackendInterface
+from lmcache.v1.server.utils import LMSMemoryObj
+
+logger = init_logger(__name__)
+
+
+class LMSLocalBackend(LMSBackendInterface):
+    def __init__(
+        self,
+    ) -> None:
+        self.dict: OrderedDict[CacheEngineKey, LMSMemoryObj] = OrderedDict()
+
+        self.lock = threading.Lock()
+
+        # TODO(Jiayi): please add evictor
+
+    # TODO
+    def list_keys(self) -> List[CacheEngineKey]:
+        with self.lock:
+            return list(self.dict.keys())
+
+    def contains(
+        self,
+        key: CacheEngineKey,
+    ) -> bool:
+        with self.lock:
+            return key in self.dict
+
+    # TODO
+    def remove(
+        self,
+        key: CacheEngineKey,
+    ) -> None:
+        with self.lock:
+            self.dict.pop(key)
+
+    def put(
+        self,
+        client_meta: ClientMetaMessage,
+        kv_chunk_bytes: bytearray,
+    ) -> None:
+        with self.lock:
+            self.dict[client_meta.key] = LMSMemoryObj(
+                kv_chunk_bytes,
+                client_meta.length,
+                client_meta.fmt,
+                client_meta.dtype,
+                client_meta.shape,
+            )
+
+    @_lmcache_nvtx_annotate
+    def get(
+        self,
+        key: CacheEngineKey,
+    ) -> Optional[LMSMemoryObj]:
+        with self.lock:
+            return self.dict.get(key, None)
+
+    def close(self):
+        pass
+
+
+# TODO(Jiayi): please implement the remote disk backend
+# class LMSLocalDiskBackend(LMSBackendInterface):
+#    pass

lmcache/v1/server/utils.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass
+from typing import Optional
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.v1.memory_management import MemoryFormat
+
+
+# TODO(Jiayi): Maybe move the memory management in remote
+# cache server to `memory_management.py` as well.
+@dataclass
+class LMSMemoryObj:
+    data: bytearray
+    length: int
+    fmt: MemoryFormat
+    dtype: Optional[torch.dtype]
+    shape: torch.Size

lmcache/v1/standalone/__init__.py

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