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

lmcache/v1/transfer_channel/__init__.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional
+
+# First Party
+from lmcache.v1.transfer_channel.abstract import BaseTransferChannel
+
+
+# TODO(Jiayi): Refactor this function when we support more channels.
+def CreateTransferChannel(
+    channel_type: str,
+    async_mode: bool,
+    role: str,
+    buffer_ptr: int,
+    buffer_size: int,
+    align_bytes: int,
+    tp_rank: int,
+    peer_init_url: str,
+    device: Optional[str] = None,
+    **kwargs,
+) -> BaseTransferChannel:
+    """
+    Create a transfer channel based on the specified channel type.
+    Supports "nixl" and "mock_memory" channel types.
+    If nixl is not available, automatically falls back to mock_memory
+    which is a mock implementation for testing purposes.
+
+    :param channel_type: Type of the transfer channel (e.g., "nixl", "mock_memory").
+    :param async_mode: Whether to operate in asynchronous mode.
+    :param role: Role of the channel (e.g., "both", "sender" or "receiver").
+    :param buffer_ptr: Pointer to the pre-allocated buffer.
+    :param buffer_size: Size of the pre-allocated buffer in bytes.
+    :param align_bytes: Alignment requirement in bytes.
+    :param tp_rank: Tensor parallel rank of the current process.
+    :param peer_init_url: Initialization URL for the peer.
+    :kwargs: Additional keyword arguments specific to the channel type.
+
+    :return: An instance of the specified transfer channel.
+    """
+
+    assert channel_type in ["nixl", "mock_memory"], (
+        f"Unsupported channel type: {channel_type}"
+    )
+
+    if channel_type == "nixl":
+        # First Party
+        from lmcache.v1.transfer_channel.nixl_channel import NixlChannel
+
+        assert "backends" in kwargs, (
+            "`backends` must be provided to create nixl transfer channel."
+        )
+        transfer_channel = NixlChannel(
+            async_mode=async_mode,
+            role=role,
+            buffer_ptr=buffer_ptr,
+            buffer_size=buffer_size,
+            align_bytes=align_bytes,
+            tp_rank=tp_rank,
+            peer_init_url=peer_init_url,
+            device=device,
+            **kwargs,
+        )
+        return transfer_channel
+
+    if channel_type == "mock_memory":
+        # First Party
+        from lmcache.v1.transfer_channel.mock_memory_channel import (
+            MockMemoryChannel,
+        )
+
+        mock_memory_channel: BaseTransferChannel = MockMemoryChannel(
+            async_mode=async_mode,
+            role=role,
+            buffer_ptr=buffer_ptr,
+            buffer_size=buffer_size,
+            align_bytes=align_bytes,
+            tp_rank=tp_rank,
+            peer_init_url=peer_init_url,
+            **kwargs,
+        )
+        return mock_memory_channel
+
+    raise ValueError(f"Unsupported channel type: {channel_type}")

lmcache/v1/transfer_channel/abstract.py

@@ -0,0 +1,285 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Union
+import abc
+
+# Third Party
+import msgspec
+import zmq
+
+# First Party
+from lmcache.v1.memory_management import MemoryObj
+from lmcache.v1.transfer_channel.transfer_utils import (
+    InitSideMsgBase,
+    InitSideRetMsgBase,
+    P2PInitSideMsg,
+    P2PInitSideRetMsg,
+    SideMsg,
+)
+
+
+class BaseTransferChannel(metaclass=abc.ABCMeta):
+    ### Initialization-related functions ###
+    @abc.abstractmethod
+    def lazy_init_peer_connection(
+        self,
+        local_id: str,
+        peer_id: str,
+        peer_init_url: str,
+        init_side_msg: Optional[InitSideMsgBase] = None,
+    ) -> Optional[InitSideRetMsgBase]:
+        """
+        Lazily initialize the connection to a peer.
+
+        :param local_id: The ID of itself.
+        peer_id: The ID of the peer to connect to.
+        peer_init_url: The URL used to initialize the connection.
+        init_side_msg: An optional side message to be sent to the peer
+        during initialization.
+
+        :return: An optional side message received from the peer.
+        """
+
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    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]:
+        """
+        Async version of `lazy_init_peer_connection`.
+        """
+
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def remote_xfer_handler_exists(self, receiver_or_sender_id: str) -> bool:
+        """
+        Check if the remote transfer handler exists.
+
+        :param receiver_or_sender_id: The ID of the receiver or sender.
+
+        :return: True if the remote transfer handler exists, False otherwise.
+        """
+        raise NotImplementedError
+
+    def handle_init_side_msg(
+        self,
+        req: InitSideMsgBase,
+    ) -> InitSideRetMsgBase:
+        """
+        Handle side messages during initialization.
+
+        :param req: The initialization-related side message
+        received from the peer.
+
+        :return: A side message to be sent back to the peer.
+        """
+        if isinstance(req, P2PInitSideMsg):
+            assert hasattr(self, "peer_lookup_url"), (
+                "P2PInitSideMsg requires `peer_lookup_url` attribute."
+            )
+            return P2PInitSideRetMsg(
+                peer_lookup_url=self.peer_lookup_url,
+            )
+        else:
+            raise ValueError(f"Unsupported InitSideMsg type: {type(req)}")
+
+    def send_init_side_msg(
+        self,
+        init_tmp_socket: zmq.Socket,
+        init_side_msg: InitSideMsgBase,
+    ) -> InitSideRetMsgBase:
+        """
+        Send side messages during initialization.
+
+        :param socket: The ZMQ socket used for sending the message.
+        :param init_side_msg: The initialization-related side message
+        to be sent to the peer.
+
+        :return: A side message received from the peer.
+        """
+        init_msg_bytes = msgspec.msgpack.encode(init_side_msg)
+        init_tmp_socket.send(init_msg_bytes)
+
+        init_ret_msg_bytes = init_tmp_socket.recv()
+        init_ret_msg = msgspec.msgpack.decode(
+            init_ret_msg_bytes,
+            type=SideMsg,
+        )
+
+        return init_ret_msg
+
+    async def async_send_init_side_msg(
+        self,
+        init_tmp_socket: zmq.Socket,
+        init_side_msg: InitSideMsgBase,
+    ) -> InitSideRetMsgBase:
+        """
+        Async version of send_init_side_msg.
+        """
+        init_msg_bytes = msgspec.msgpack.encode(init_side_msg)
+        await init_tmp_socket.send(init_msg_bytes)
+
+        init_ret_msg_bytes = await init_tmp_socket.recv()
+        init_ret_msg = msgspec.msgpack.decode(
+            init_ret_msg_bytes,
+            type=SideMsg,
+        )
+
+        return init_ret_msg
+
+    ### Utility functions ###
+    @abc.abstractmethod
+    def get_local_mem_indices(
+        self, objects: Union[list[bytes], list[MemoryObj]]
+    ) -> list[int]:
+        """
+        Get the memory indices of objects.
+
+        :param objects: A list of bytes or MemoryObj to be checked.
+
+        :return: The memory indices of the objects.
+        """
+        raise NotImplementedError
+
+    ### Send and Recv must be called in pair ###
+    @abc.abstractmethod
+    def batched_send(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Send a batch of data through the channel.
+
+        :param objects: A list of bytes or MemoryObj to be sent.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def batched_recv(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Receive a batch of data through the channel.
+
+        :param buffer: A list of bytes or MemoryObj to store the received data.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def async_batched_send(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Async send a batch of data through the channel.
+
+        :param objects: A list of bytes or MemoryObj to be sent.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def async_batched_recv(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Async receive a batch of data through the channel.
+
+        :param buffers: A list of bytes or MemoryObj to store the received data.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    ### Read and Write only need to be called on one side ###
+    @abc.abstractmethod
+    def batched_write(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Write a batch of data through the channel.
+
+        :param objects: A list of bytes or MemoryObj to be written.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def batched_read(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Read a batch of data through the channel.
+
+        :param buffers: A list of bytes or MemoryObj to store the read data.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def async_batched_write(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Async write a batch of data through the channel.
+
+        :param objects: A list of bytes or MemoryObj to be written.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def async_batched_read(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """
+        Async read a batch of data through the channel.
+
+        :param buffers: A list of bytes or MemoryObj to store the read data.
+        :param transfer_spec: Additional specifications for the transfer.
+
+        :return: Number of successfully transferred objects.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """
+        Close the transfer channel and release any resources.
+        """
+        raise NotImplementedError

lmcache/v1/transfer_channel/mock_memory_channel.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Union
+import asyncio
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.memory_management import MemoryObj
+from lmcache.v1.transfer_channel.py_socket_channel import PySocketChannel
+
+logger = init_logger(__name__)
+
+
+# Data plane: Global shared data store for simulating data transfer
+# within a single process. This is a process-local Python dict,
+# NOT shared across processes or machines.
+_GLOBAL_DATA_STORE: dict[str, dict[int, torch.Tensor]] = {}
+
+
+class MockMemoryChannel(PySocketChannel):
+    """
+    Mock memory-based transfer channel for single-process testing.
+
+    Control plane: Inherits ZMQ-based handshake from PySocketChannel.
+    Data plane: Uses global dict to simulate tensor transfer.
+
+    Only works within a single process. For production use NixlChannel.
+    """
+
+    def __init__(
+        self,
+        async_mode: bool = False,
+        **kwargs,
+    ):
+        super().__init__(async_mode=async_mode, **kwargs)
+
+    ############################################################
+    # Data plane: Override hook to initialize global data store
+    ############################################################
+    def _on_peer_connected(self, peer_url: str):
+        """Initialize data store when a peer connects"""
+        if peer_url not in _GLOBAL_DATA_STORE:
+            _GLOBAL_DATA_STORE[peer_url] = {}
+
+    ############################################################
+    # Data plane: Utility functions
+    ############################################################
+    def get_local_mem_indices(
+        self, objects: Union[list[bytes], list[MemoryObj]]
+    ) -> list[int]:
+        local_indices = []
+        if isinstance(objects[0], MemoryObj):
+            for mem_obj in objects:
+                assert isinstance(mem_obj, MemoryObj)
+                local_indices.append(mem_obj.meta.address)
+        elif isinstance(objects[0], bytes):
+            raise NotImplementedError(
+                "Sending raw bytes is not supported in MockMemoryChannel"
+            )
+        return local_indices
+
+    ############################################################
+    # Data plane: Send/Recv functions (not implemented)
+    ############################################################
+    def batched_send(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    def batched_recv(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    async def async_batched_send(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    async def async_batched_recv(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    ############################################################
+    # Data plane: Read/Write functions (global dict simulation)
+    ############################################################
+    def remote_xfer_handler_exists(self, receiver_or_sender_id: str) -> bool:
+        return True
+
+    def batched_write(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        raise NotImplementedError("Sync mode not supported in MockMemoryChannel")
+
+    def batched_read(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        raise NotImplementedError("Sync mode not supported in MockMemoryChannel")
+
+    async def async_batched_write(
+        self,
+        objects: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """Write tensors to global dict to simulate data transfer"""
+        assert transfer_spec is not None
+
+        receiver_id = transfer_spec["receiver_id"]
+        remote_indexes = transfer_spec["remote_indexes"]
+
+        for obj, remote_idx in zip(objects, remote_indexes, strict=True):
+            if isinstance(obj, MemoryObj) and obj.tensor is not None:
+                if receiver_id not in _GLOBAL_DATA_STORE:
+                    _GLOBAL_DATA_STORE[receiver_id] = {}
+                _GLOBAL_DATA_STORE[receiver_id][remote_idx] = obj.tensor.clone()
+
+        await asyncio.sleep(0.001)
+        return len(objects)
+
+    async def async_batched_read(
+        self,
+        buffers: Union[list[bytes], list[MemoryObj]],
+        transfer_spec: Optional[dict] = None,
+    ) -> int:
+        """Read tensors from global dict to simulate data transfer"""
+        assert transfer_spec is not None
+
+        sender_id = transfer_spec["sender_id"]
+        remote_indexes = transfer_spec["remote_indexes"]
+
+        for buf, remote_idx in zip(buffers, remote_indexes, strict=True):
+            if isinstance(buf, MemoryObj) and buf.tensor is not None:
+                if (
+                    sender_id in _GLOBAL_DATA_STORE
+                    and remote_idx in _GLOBAL_DATA_STORE[sender_id]
+                ):
+                    buf.tensor.copy_(_GLOBAL_DATA_STORE[sender_id][remote_idx])
+
+        await asyncio.sleep(0.001)
+        return len(buffers)