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

lmcache/v1/storage_backend/connector/infinistore_connector.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from functools import reduce
+from typing import List, Optional, Union, no_type_check
+import asyncio
+import ctypes
+import operator
+
+# Third Party
+import infinistore
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.utils import CacheEngineKey
+from lmcache.v1.memory_management import MemoryObj
+
+# reuse
+from lmcache.v1.protocol import RemoteMetadata
+from lmcache.v1.storage_backend.connector.base_connector import RemoteConnector
+from lmcache.v1.storage_backend.local_cpu_backend import LocalCPUBackend
+
+logger = init_logger(__name__)
+
+MAX_BUFFER_SIZE = 40 << 20  # 40MB
+METADATA_BYTES_LEN = 28
+MAX_BUFFER_CNT = 16
+
+
+def _get_ptr(mv: Union[bytearray, memoryview]) -> int:
+    return ctypes.addressof(ctypes.c_char.from_buffer(mv))
+
+
+class InfinistoreConnector(RemoteConnector):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        dev_name: str,
+        link_type: str,
+        loop: asyncio.AbstractEventLoop,
+        memory_allocator: LocalCPUBackend,
+    ):
+        # initialize base class, which includes some common attributes
+        super().__init__(memory_allocator.config, memory_allocator.metadata)
+
+        config = infinistore.ClientConfig(
+            host_addr=host,
+            service_port=port,
+            log_level="info",
+            connection_type=infinistore.TYPE_RDMA,
+            ib_port=1,
+            link_type=link_type,
+            dev_name=dev_name,
+        )
+
+        self.rdma_conn = infinistore.InfinityConnection(config)
+
+        self.loop = loop
+        self.rdma_conn.connect()
+
+        self.send_buffers = []
+        self.recv_buffers = []
+        self.send_queue: asyncio.Queue[int] = asyncio.Queue(maxsize=MAX_BUFFER_CNT)
+        self.recv_queue: asyncio.Queue[int] = asyncio.Queue(maxsize=MAX_BUFFER_CNT)
+
+        self.buffer_size = MAX_BUFFER_SIZE
+        self.memory_allocator = memory_allocator
+
+        for i in range(MAX_BUFFER_CNT):
+            send_buffer = bytearray(self.buffer_size)
+            self.rdma_conn.register_mr(_get_ptr(send_buffer), self.buffer_size)
+            self.send_buffers.append(send_buffer)
+            self.send_queue.put_nowait(i)
+
+            recv_buffer = bytearray(self.buffer_size)
+            self.rdma_conn.register_mr(_get_ptr(recv_buffer), self.buffer_size)
+            self.recv_buffers.append(recv_buffer)
+            self.recv_queue.put_nowait(i)
+
+    async def exists(self, key: CacheEngineKey) -> bool:
+        def blocking_io():
+            return self.rdma_conn.check_exist(key.to_string())
+
+        return await self.loop.run_in_executor(None, blocking_io)
+
+    def exists_sync(self, key: CacheEngineKey) -> bool:
+        return self.rdma_conn.check_exist(key.to_string())
+
+    async def get(self, key: CacheEngineKey) -> Optional[MemoryObj]:
+        key_str = key.to_string()
+
+        buf_idx = await self.recv_queue.get()
+        buffer = self.recv_buffers[buf_idx]
+        try:
+            await self.rdma_conn.rdma_read_cache_async(
+                [(key_str, 0)], self.buffer_size, _get_ptr(buffer)
+            )
+        except Exception as e:
+            logger.warning(f"get failed: {e}")
+            self.recv_queue.put_nowait(buf_idx)
+            return None
+
+        metadata = RemoteMetadata.deserialize(buffer)
+
+        num_elements = reduce(operator.mul, metadata.shapes[0])
+        assert len(metadata.dtypes) == 1
+        temp_tensor = torch.frombuffer(
+            buffer,
+            dtype=metadata.dtypes[0],
+            offset=METADATA_BYTES_LEN,
+            count=num_elements,
+        ).reshape(metadata.shapes[0])
+
+        memory_obj = self.memory_allocator.allocate(
+            metadata.shapes[0],
+            metadata.dtypes[0],
+            metadata.fmt,
+        )
+
+        assert memory_obj is not None
+        assert memory_obj.tensor is not None
+
+        # deep copy to pinned memory
+        # and hot cache will reference this memory obj
+        memory_obj.tensor.copy_(temp_tensor)
+
+        logger.debug(f"get key: {key_str} done, {memory_obj.get_shape()}")
+        self.recv_queue.put_nowait(buf_idx)
+
+        return memory_obj
+
+    async def put(self, key: CacheEngineKey, memory_obj: MemoryObj):
+        key_str = key.to_string()
+
+        kv_bytes = memory_obj.byte_array
+        kv_shapes = memory_obj.get_shapes()
+        kv_dtypes = memory_obj.get_dtypes()
+        memory_format = memory_obj.get_memory_format()
+
+        buf_idx = await self.send_queue.get()
+        buffer = self.send_buffers[buf_idx]
+
+        RemoteMetadata(
+            len(kv_bytes), kv_shapes, kv_dtypes, memory_format
+        ).serialize_into(buffer)
+
+        buffer[METADATA_BYTES_LEN : METADATA_BYTES_LEN + len(kv_bytes)] = kv_bytes
+
+        size = memory_obj.get_physical_size()
+
+        if size + METADATA_BYTES_LEN > self.buffer_size:
+            raise ValueError(
+                f"Value size ({size + METADATA_BYTES_LEN} bytes)"
+                f"exceeds the maximum allowed size"
+                f"({self.buffer_size} bytes). Please decrease chunk_size."
+            )
+        try:
+            await self.rdma_conn.rdma_write_cache_async(
+                [(key_str, 0)], METADATA_BYTES_LEN + size, _get_ptr(buffer)
+            )
+        except Exception as e:
+            logger.warning(f"exception happens in rdma_write_cache_async kv_bytes {e}")
+            return
+        finally:
+            self.send_queue.put_nowait(buf_idx)
+
+        logger.debug(f"put key: {key.to_string()}, {memory_obj.get_shape()}")
+
+    # TODO
+    @no_type_check
+    async def list(self) -> List[str]:
+        pass
+
+    async def close(self):
+        self.rdma_conn.close()
+        logger.info("Closed the infinistore connection")

lmcache/v1/storage_backend/connector/instrumented_connector.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import List, Optional
+import time
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.observability import LMCStatsMonitor
+from lmcache.utils import CacheEngineKey
+from lmcache.v1.memory_management import MemoryObj
+from lmcache.v1.storage_backend.connector.base_connector import RemoteConnector
+
+logger = init_logger(__name__)
+
+
+class InstrumentedRemoteConnector(RemoteConnector):
+    """
+    A connector that instruments the underlying connector with
+    metrics collection and logging capabilities.
+    """
+
+    def __init__(self, connector: RemoteConnector):
+        self._connector = connector
+        self._stats_monitor = LMCStatsMonitor.GetOrCreate()
+        self.name = self.__repr__()
+
+    async def put(self, key: CacheEngineKey, memory_obj: MemoryObj) -> None:
+        obj_size = memory_obj.get_size()
+        begin = time.perf_counter()
+        try:
+            await self._connector.put(key, memory_obj)
+        finally:
+            # Ensure reference count is decreased even if exception occurs
+            memory_obj.ref_count_down()
+
+        end = time.perf_counter()
+        self._stats_monitor.update_interval_remote_time_to_put((end - begin) * 1000)
+        self._stats_monitor.update_interval_remote_write_metrics(obj_size)
+        logger.debug(
+            "[%s]Bytes offloaded: %.3f MBytes in %.3f ms",
+            self.name,
+            obj_size / 1e6,
+            (end - begin) * 1000,
+        )
+
+    async def get(self, key: CacheEngineKey) -> Optional[MemoryObj]:
+        begin = time.perf_counter()
+        memory_obj = await self._connector.get(key)
+        end = time.perf_counter()
+        duration = end - begin
+
+        retrieve_stats = self._stats_monitor.get_current_retrieve_stats()
+        if (
+            retrieve_stats is not None
+            and "remote_backend_individual_get_stats" in retrieve_stats.detailed_metrics
+        ):
+            retrieve_stats.detailed_metrics["remote_backend_individual_get_stats"][
+                key
+            ] = {"instrumented_connector_get_time": duration}
+
+        if memory_obj is not None:
+            obj_size = memory_obj.get_size()
+            self._stats_monitor.update_interval_remote_read_metrics(obj_size)
+            logger.debug(
+                "[%s]Bytes loaded: %.3f MBytes in %.3f ms",
+                self.name,
+                obj_size / 1e6,
+                duration * 1000,
+            )
+        return memory_obj
+
+    # Delegate all other methods to the underlying connector
+    async def exists(self, key: CacheEngineKey) -> bool:
+        return await self._connector.exists(key)
+
+    def exists_sync(self, key: CacheEngineKey) -> bool:
+        return self._connector.exists_sync(key)
+
+    async def list(self) -> List[str]:
+        return await self._connector.list()
+
+    async def close(self) -> None:
+        await self._connector.close()
+
+    def getWrappedConnector(self) -> RemoteConnector:
+        return self._connector
+
+    def support_ping(self) -> bool:
+        return self._connector.support_ping()
+
+    async def ping(self) -> int:
+        return await self._connector.ping()
+
+    def support_batched_put(self) -> bool:
+        return self._connector.support_batched_put()
+
+    def support_batched_get(self) -> bool:
+        return self._connector.support_batched_get()
+
+    def support_batched_async_contains(self) -> bool:
+        return self._connector.support_batched_async_contains()
+
+    async def batched_async_contains(
+        self,
+        lookup_id: str,
+        keys: List[CacheEngineKey],
+        pin: bool = False,
+    ) -> int:
+        return await self._connector.batched_async_contains(lookup_id, keys, pin)
+
+    def support_batched_get_non_blocking(self) -> bool:
+        return self._connector.support_batched_get_non_blocking()
+
+    async def batched_get_non_blocking(
+        self,
+        lookup_id: str,
+        keys: List[CacheEngineKey],
+    ) -> List[MemoryObj]:
+        begin = time.perf_counter()
+        memory_objs = await self._connector.batched_get_non_blocking(lookup_id, keys)
+        end = time.perf_counter()
+        duration = end - begin
+
+        total_size = sum(
+            memory_obj.get_size()
+            for memory_obj in memory_objs
+            if memory_obj is not None
+        )
+        if total_size > 0:
+            self._stats_monitor.update_interval_remote_read_metrics(total_size)
+            logger.debug(
+                "[%s]Bytes loaded: %.3f MBytes in %.3f ms",
+                self.name,
+                total_size / 1e6,
+                duration * 1000,
+            )
+        return memory_objs
+
+    async def batched_get(
+        self, keys: List[CacheEngineKey]
+    ) -> List[Optional[MemoryObj]]:
+        begin = time.perf_counter()
+        memory_objs = await self._connector.batched_get(keys)
+        end = time.perf_counter()
+        duration = end - begin
+        self._stats_monitor.update_interval_remote_time_to_get(duration * 1000)
+
+        retrieve_stats = self._stats_monitor.get_current_retrieve_stats()
+        if retrieve_stats is not None:
+            retrieve_stats.detailed_metrics[
+                "instrumented_connector_batched_get_time"
+            ] = (
+                retrieve_stats.detailed_metrics.get(
+                    "instrumented_connector_batched_get_time", 0.0
+                )
+                + duration
+            )
+
+        total_size = sum(
+            memory_obj.get_size()
+            for memory_obj in memory_objs
+            if memory_obj is not None
+        )
+        if total_size > 0:
+            self._stats_monitor.update_interval_remote_read_metrics(total_size)
+            logger.debug(
+                "[%s]Bytes loaded: %.3f MBytes in %.3f ms",
+                self.name,
+                total_size / 1e6,
+                duration * 1000,
+            )
+        return memory_objs
+
+    async def batched_put(
+        self, keys: List[CacheEngineKey], memory_objs: List[MemoryObj]
+    ):
+        total_size = sum(
+            memory_obj.get_size()
+            for memory_obj in memory_objs
+            if memory_obj is not None
+        )
+        begin = time.perf_counter()
+        try:
+            await self._connector.batched_put(keys, memory_objs)
+        except Exception as e:
+            logger.warning(f"batched put error: {e}")
+        finally:
+            for memory_obj in memory_objs:
+                memory_obj.ref_count_down()
+
+        end = time.perf_counter()
+        self._stats_monitor.update_interval_remote_time_to_put((end - begin) * 1000)
+        self._stats_monitor.update_interval_remote_write_metrics(total_size)
+        logger.debug(
+            "[%s]Bytes offloaded: %.3f MBytes in %.3f ms",
+            self.name,
+            total_size / 1e6,
+            (end - begin) * 1000,
+        )
+
+    def remove_sync(self, key: CacheEngineKey) -> bool:
+        return self._connector.remove_sync(key)
+
+    def batched_contains(self, keys: List[CacheEngineKey]) -> int:
+        return self._connector.batched_contains(keys)
+
+    def support_batched_contains(self) -> bool:
+        return self._connector.support_batched_contains()
+
+    def reshape_partial_chunk(
+        self, memory_obj: MemoryObj, bytes_read: int
+    ) -> MemoryObj:
+        return self._connector.reshape_partial_chunk(memory_obj, bytes_read)
+
+    def post_init(self):
+        return self._connector.post_init()
+
+    def __repr__(self) -> str:
+        return f"InstrumentedRemoteConnector({self._connector})"

lmcache/v1/storage_backend/connector/lm_adapter.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.storage_backend.connector import (
+    ConnectorAdapter,
+    ConnectorContext,
+    parse_remote_url,
+)
+from lmcache.v1.storage_backend.connector.base_connector import RemoteConnector
+
+logger = init_logger(__name__)
+
+
+class LMServerConnectorAdapter(ConnectorAdapter):
+    """Adapter for LM Server connectors."""
+
+    def __init__(self) -> None:
+        super().__init__("lm://")
+
+    def create_connector(self, context: ConnectorContext) -> RemoteConnector:
+        # Local
+        from .lm_connector import LMCServerConnector
+
+        logger.info(f"Creating LM Server connector for URL: {context.url}")
+        parse_url = parse_remote_url(context.url)
+        return LMCServerConnector(
+            host=parse_url.host,
+            port=parse_url.port,
+            loop=context.loop,
+            local_cpu_backend=context.local_cpu_backend,
+        )

lmcache/v1/storage_backend/connector/lm_connector.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import List, Optional, no_type_check
+import asyncio
+import socket
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.utils import CacheEngineKey, _lmcache_nvtx_annotate
+from lmcache.v1.memory_management import MemoryFormat, MemoryObj
+from lmcache.v1.protocol import (
+    ClientCommand,
+    ClientMetaMessage,
+    ServerMetaMessage,
+    ServerReturnCode,
+)
+from lmcache.v1.storage_backend.connector.base_connector import RemoteConnector
+from lmcache.v1.storage_backend.local_cpu_backend import LocalCPUBackend
+
+logger = init_logger(__name__)
+
+
+# TODO: performance optimization for this class, consider using C/C++/Rust
+# for communication + deserialization
+class LMCServerConnector(RemoteConnector):
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        loop: asyncio.AbstractEventLoop,
+        local_cpu_backend: LocalCPUBackend,
+    ):
+        # NOTE(Jiayi): According to Python documentation:
+        # https://docs.python.org/3/library/asyncio-eventloop.html
+        # In general, protocol implementations that use transport-based APIs
+        # such as loop.create_connection() and loop.create_server() are faster
+        # than implementations that work with sockets.
+        # However, we use socket here as we need to use the socket.recv_into()
+        # to reduce memory copy.
+
+        # initialize base class, which includes some common attributes
+        super().__init__(local_cpu_backend.config, local_cpu_backend.metadata)
+
+        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.client_socket.connect((host, port))
+        # loop.sock_recv_into(sock, buf)
+
+        self.loop = loop
+        self.local_cpu_backend = local_cpu_backend
+
+        self.async_socket_lock = asyncio.Lock()
+
+    # TODO(Jiayi): This should be an async function
+    def receive_all(self, meta: ServerMetaMessage) -> Optional[MemoryObj]:
+        received = 0
+        n = meta.length
+
+        # TODO(Jiayi): Format will be used once we support
+        # compressed memory format
+        memory_obj = self.local_cpu_backend.allocate(
+            meta.shape,
+            meta.dtype,
+            meta.fmt,
+        )
+        if memory_obj is None:
+            logger.warning("Failed to allocate memory during remote receive")
+            return None
+
+        buffer = memory_obj.byte_array
+        view = memoryview(buffer)
+
+        while received < n:
+            num_bytes = self.client_socket.recv_into(view[received:], n - received)
+            if num_bytes == 0:
+                return None
+            received += num_bytes
+
+        return memory_obj
+
+    async def exists(self, key: CacheEngineKey) -> bool:
+        # logger.debug("Call to exists()!")
+
+        async with self.async_socket_lock:
+            self.client_socket.sendall(
+                ClientMetaMessage(
+                    ClientCommand.EXIST,
+                    key,
+                    0,
+                    MemoryFormat(1),
+                    torch.float16,
+                    torch.Size([0, 0, 0, 0]),
+                ).serialize()
+            )
+
+            response = self.client_socket.recv(ServerMetaMessage.packlength())
+
+        return ServerMetaMessage.deserialize(response).code == ServerReturnCode.SUCCESS
+
+    def exists_sync(self, key: CacheEngineKey) -> bool:
+        future = asyncio.run_coroutine_threadsafe(self.exists(key), self.loop)
+        try:
+            res = future.result()
+            return res
+        except Exception as e:
+            logger.warning(f"lm connector failed in exists: {e}")
+            return False
+
+    async def put(
+        self,
+        key: CacheEngineKey,
+        memory_obj: MemoryObj,
+    ):
+        # logger.debug("Async call to put()!")
+
+        kv_bytes = memory_obj.byte_array
+        kv_shape = memory_obj.get_shape()
+        kv_dtype = memory_obj.get_dtype()
+        memory_format = memory_obj.get_memory_format()
+
+        async with self.async_socket_lock:
+            await self.loop.sock_sendall(
+                self.client_socket,
+                ClientMetaMessage(
+                    ClientCommand.PUT,
+                    key,
+                    len(kv_bytes),
+                    memory_format,
+                    kv_dtype,
+                    kv_shape,
+                ).serialize(),
+            )
+
+            await self.loop.sock_sendall(self.client_socket, kv_bytes)
+
+    # TODO(Jiayi): This should be an async function
+    @_lmcache_nvtx_annotate
+    async def get(self, key: CacheEngineKey) -> Optional[MemoryObj]:
+        # NOTE(Jiayi): Not using any await in the following as
+        # we don't want to yield control to other tasks which could
+        # sacrifice the performance loading to trade the performance of
+        # saving
+        async with self.async_socket_lock:
+            self.client_socket.sendall(
+                ClientMetaMessage(
+                    ClientCommand.GET,
+                    key,
+                    0,
+                    MemoryFormat(1),
+                    torch.float16,
+                    torch.Size([0, 0, 0, 0]),
+                ).serialize()
+            )
+
+            data = self.client_socket.recv(ServerMetaMessage.packlength())
+
+        meta = ServerMetaMessage.deserialize(data)
+        if meta.code != ServerReturnCode.SUCCESS:
+            return None
+
+        async with self.async_socket_lock:
+            memory_obj = self.receive_all(meta)
+
+        return memory_obj
+
+    # TODO
+    @no_type_check
+    async def list(self) -> List[str]:
+        pass
+
+    async def close(self):
+        async with self.async_socket_lock:
+            self.client_socket.close()
+        logger.info("Closed the lmserver connection")

lmcache/v1/storage_backend/connector/mock_adapter.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from urllib.parse import parse_qs, urlparse
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.storage_backend.connector import (
+    ConnectorAdapter,
+    ConnectorContext,
+)
+from lmcache.v1.storage_backend.connector.base_connector import RemoteConnector
+
+logger = init_logger(__name__)
+
+
+class MockConnectorAdapter(ConnectorAdapter):
+    """Adapter for Mock Connector"""
+
+    def __init__(self) -> None:
+        super().__init__("mock://")
+
+    def create_connector(self, context: ConnectorContext) -> RemoteConnector:
+        # Local import to avoid circular dependencies
+        # Local
+        from .mock_connector import MockConnector
+
+        logger.info(f"Creating Mock connector for URL: {context.url}")
+
+        parsed = urlparse(context.url)
+        # capacity is provided as the netloc in URLs like: mock://100/?...
+        if not parsed.netloc:
+            raise ValueError(
+                "mock connector requires capacity in GB as netloc, e.g. mock://100/?..."
+            )
+        try:
+            capacity_gb = int(parsed.netloc)
+        except ValueError as e:
+            raise ValueError(
+                f"Invalid capacity '{parsed.netloc}' for",
+                " mock connector; must be an integer (GB).",
+            ) from e
+
+        params = parse_qs(parsed.query) if parsed.query else {}
+        # Defaults
+        peeking_latency_ms = float(params.get("peeking_latency", ["1"])[0])
+        read_throughput_gbps = float(params.get("read_throughput", ["2"])[0])
+        write_throughput_gbps = float(params.get("write_throughput", ["2"])[0])
+
+        return MockConnector(
+            url=context.url,
+            loop=context.loop,
+            local_cpu_backend=context.local_cpu_backend,
+            capacity=capacity_gb,
+            read_throughput=read_throughput_gbps,
+            write_throughput=write_throughput_gbps,
+            peeking_latency=peeking_latency_ms,
+        )