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

lmcache/v1/cache_controller/controller_manager.py

@@ -0,0 +1,535 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional, Union
+import asyncio
+import json
+import threading
+import time
+
+# Third Party
+import msgspec
+import zmq
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.controllers import KVController, RegistrationController
+from lmcache.v1.cache_controller.executor import LMCacheClusterExecutor
+from lmcache.v1.cache_controller.observability import (
+    PrometheusLogger,
+    SocketMetricsContext,
+    SocketType,
+)
+from lmcache.v1.rpc_utils import (
+    get_ip,
+    get_zmq_context,
+    get_zmq_socket,
+)
+
+from lmcache.v1.cache_controller.message import (  # isort: skip
+    BatchedKVOperationMsg,
+    BatchedP2PLookupMsg,
+    CheckFinishMsg,
+    ClearMsg,
+    CompressMsg,
+    DecompressMsg,
+    DeRegisterMsg,
+    ErrorMsg,
+    FullSyncBatchMsg,
+    FullSyncEndMsg,
+    FullSyncStartMsg,
+    FullSyncStatusMsg,
+    HealthMsg,
+    HeartbeatMsg,
+    LookupMsg,
+    MoveMsg,
+    Msg,
+    MsgBase,
+    OrchMsg,
+    OrchRetMsg,
+    PinMsg,
+    QueryInstMsg,
+    QueryWorkerInfoMsg,
+    RegisterMsg,
+    WorkerMsg,
+    WorkerReqMsg,
+    WorkerReqRetMsg,
+)
+
+logger = init_logger(__name__)
+
+# TODO(Jiayi): Need to align the message types. For example,
+# a controller should take in an control message and return
+# a control message.
+
+
+class LMCacheControllerManager:
+    def __init__(
+        self,
+        controller_urls: dict[str, str],
+        health_check_interval: int,
+        lmcache_worker_timeout: int,
+        full_sync_completion_threshold: float = 0.8,
+        full_sync_timeout_s: float = 300.0,
+    ):
+        # Initialize stats logger
+        prometheus_labels = {
+            "role": "controller",
+        }
+        PrometheusLogger.GetOrCreate(prometheus_labels)
+        self.zmq_context = get_zmq_context()
+        self.controller_urls = controller_urls
+        # TODO(Jiayi): We might need multiple sockets if there are more
+        # controllers. For now, we use a single socket to receive messages
+        # for all controllers.
+        # Similarly we might need more sockets to handle different control
+        # messages. For now, we use one socket to handle all control messages.
+
+        # TODO(Jiayi): Another thing is that we might need to decoupe the
+        # interactions among `handle_worker_message`, `handle_control_message`
+        # and `handle_orchestration_message`. For example, in
+        # `handle_orchestration_message`, we might need to call
+        # `issue_control_message`. This will make the system less concurrent.
+
+        # Micro controllers
+        self.controller_pull_socket = get_zmq_socket(
+            self.zmq_context,
+            self.controller_urls["pull"],
+            protocol="tcp",
+            role=zmq.PULL,  # type: ignore[attr-defined]
+            bind_or_connect="bind",
+        )
+
+        if self.controller_urls["reply"] is not None:
+            self.controller_reply_socket = get_zmq_socket(
+                self.zmq_context,
+                self.controller_urls["reply"],
+                protocol="tcp",
+                role=zmq.ROUTER,  # type: ignore[attr-defined]
+                bind_or_connect="bind",
+            )
+
+        # Dedicated heartbeat socket to avoid blocking from other requests
+        if self.controller_urls.get("heartbeat") is not None:
+            self.controller_heartbeat_socket = get_zmq_socket(
+                self.zmq_context,
+                self.controller_urls["heartbeat"],
+                protocol="tcp",
+                role=zmq.ROUTER,  # type: ignore[attr-defined]
+                bind_or_connect="bind",
+            )
+        else:
+            self.controller_heartbeat_socket = None
+        self.reg_controller = RegistrationController()
+        self.kv_controller = KVController(
+            registry=self.reg_controller.registry,
+            full_sync_completion_threshold=full_sync_completion_threshold,
+            full_sync_timeout_s=full_sync_timeout_s,
+        )
+
+        # Cluster executor
+        self.cluster_executor = LMCacheClusterExecutor(
+            reg_controller=self.reg_controller,
+        )
+
+        # post initialization of controllers
+        self.kv_controller.post_init(
+            reg_controller=self.reg_controller,
+            cluster_executor=self.cluster_executor,
+        )
+        self.reg_controller.post_init(
+            kv_controller=self.kv_controller,
+            cluster_executor=self.cluster_executor,
+        )
+        self.health_check_interval = health_check_interval
+        self.lmcache_worker_timeout = lmcache_worker_timeout
+
+        if self.health_check_interval > 0:
+            logger.info(
+                "Start health check thread, interval: %s", self.health_check_interval
+            )
+            self.loop = asyncio.new_event_loop()
+            self.thread = threading.Thread(
+                target=self.loop.run_forever,
+                daemon=True,
+                name="controller-health-thread",
+            )
+            self.thread.start()
+            asyncio.run_coroutine_threadsafe(self.health_check(), self.loop)
+
+        # Setup socket message count metrics
+        self._setup_socket_metrics()
+
+    async def handle_worker_message(self, msg: WorkerMsg) -> None:
+        if isinstance(msg, RegisterMsg):
+            await self.reg_controller.register(msg)
+        elif isinstance(msg, DeRegisterMsg):
+            await self.reg_controller.deregister(msg)
+        elif isinstance(msg, BatchedKVOperationMsg):
+            await self.kv_controller.handle_batched_kv_operations(msg)
+        elif isinstance(msg, FullSyncBatchMsg):
+            await self.kv_controller.handle_full_sync_batch(msg)
+        elif isinstance(msg, FullSyncEndMsg):
+            await self.kv_controller.handle_full_sync_end(msg)
+        else:
+            logger.error(f"Unknown worker message type: {msg}")
+
+    async def handle_worker_req_message(
+        self, msg: WorkerReqMsg
+    ) -> Union[WorkerReqRetMsg, ErrorMsg]:
+        ret_msg: Union[WorkerReqRetMsg, ErrorMsg]
+        if isinstance(msg, RegisterMsg):
+            # Build extra_config with heartbeat_url if available
+            extra_config: dict[str, str] = {}
+            if self.controller_urls.get("heartbeat") is not None:
+                # Convert bind address (e.g., "0.0.0.0:8082" or "*:8082")
+                # to a connectable address using actual controller IP
+                heartbeat_bind_url = self.controller_urls["heartbeat"]
+                heartbeat_url = self._convert_bind_to_connect_url(
+                    heartbeat_bind_url, worker_ip=msg.ip
+                )
+                extra_config["heartbeat_url"] = heartbeat_url
+                logger.debug(
+                    "Returning heartbeat_url to worker: %s (bind: %s)",
+                    heartbeat_url,
+                    heartbeat_bind_url,
+                )
+            ret_msg = await self.reg_controller.register(msg, extra_config)
+        elif isinstance(msg, BatchedP2PLookupMsg):
+            ret_msg = await self.kv_controller.batched_p2p_lookup(msg)
+        elif isinstance(msg, HeartbeatMsg):
+            ret_msg = await self.reg_controller.heartbeat(msg)
+        elif isinstance(msg, FullSyncStartMsg):
+            ret_msg = await self.kv_controller.handle_full_sync_start(msg)
+        elif isinstance(msg, FullSyncStatusMsg):
+            ret_msg = await self.kv_controller.handle_full_sync_status(msg)
+        else:
+            logger.error(f"Unknown worker request message type: {msg}")
+            ret_msg = ErrorMsg(error=f"Unknown message type: {type(msg)}")
+        return ret_msg
+
+    async def handle_orchestration_message(self, msg: OrchMsg) -> OrchRetMsg:
+        if isinstance(msg, LookupMsg):
+            return await self.kv_controller.lookup(msg)
+        elif isinstance(msg, HealthMsg):
+            return await self.reg_controller.health(msg)
+        elif isinstance(msg, QueryInstMsg):
+            return await self.reg_controller.get_instance_id(msg)
+        elif isinstance(msg, ClearMsg):
+            return await self.kv_controller.clear(msg)
+        elif isinstance(msg, PinMsg):
+            return await self.kv_controller.pin(msg)
+        elif isinstance(msg, CompressMsg):
+            return await self.kv_controller.compress(msg)
+        elif isinstance(msg, DecompressMsg):
+            return await self.kv_controller.decompress(msg)
+        elif isinstance(msg, MoveMsg):
+            return await self.kv_controller.move(msg)
+        elif isinstance(msg, CheckFinishMsg):
+            # FIXME(Jiayi): This `check_finish` thing
+            # shouldn't be implemented in kv_controller.
+            return await self.kv_controller.check_finish(msg)
+        elif isinstance(msg, QueryWorkerInfoMsg):
+            return await self.reg_controller.query_worker_info(msg)
+        else:
+            logger.error(f"Unknown orchestration message type: {msg}")
+            raise RuntimeError(f"Unknown orchestration message type: {msg}")
+
+    def _setup_socket_metrics(self):
+        """Setup metrics for socket message counts."""
+        # Initialize message counters for observability
+        self.pull_socket_message_count = 0
+        self.reply_socket_message_count = 0
+
+        # Initialize active request counters
+        self.pull_socket_active_requests = 0
+        self.reply_socket_active_requests = 0
+
+        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
+        if prometheus_logger is not None:
+            prometheus_logger.pull_socket_message_count.set_function(
+                lambda: self.pull_socket_message_count
+            )
+            prometheus_logger.reply_socket_message_count.set_function(
+                lambda: self.reply_socket_message_count
+            )
+
+            # Socket queue/backlog metrics
+            prometheus_logger.pull_socket_has_pending.set_function(
+                lambda: self._check_socket_has_pending(self.controller_pull_socket)
+            )
+            if self.controller_urls["reply"] is not None:
+                prometheus_logger.reply_socket_has_pending.set_function(
+                    lambda: self._check_socket_has_pending(self.controller_reply_socket)
+                )
+
+            # Active request metrics
+            prometheus_logger.pull_socket_active_requests.set_function(
+                lambda: self.pull_socket_active_requests
+            )
+            prometheus_logger.reply_socket_active_requests.set_function(
+                lambda: self.reply_socket_active_requests
+            )
+
+    def _convert_bind_to_connect_url(
+        self, bind_url: str, worker_ip: Optional[str] = None
+    ) -> str:
+        """Convert a bind address to a connectable address.
+
+        Bind addresses like "0.0.0.0:port" or "*:port" cannot be used
+        by workers to connect. We need to replace them with the actual
+        controller IP address.
+
+        If worker_ip is provided and matches the controller's IP, use
+        127.0.0.1 for loopback connection (more reliable than external IP).
+
+        Args:
+            bind_url: The bind URL (e.g., "0.0.0.0:8082" or "*:8082")
+            worker_ip: The IP address of the requesting worker (optional)
+
+        Returns:
+            A connectable URL (e.g., "192.168.1.100:8082" or "127.0.0.1:8082")
+        """
+        if ":" not in bind_url:
+            return bind_url
+
+        host, port = bind_url.rsplit(":", 1)
+        # Replace bind-all addresses with actual IP
+        if host in ("0.0.0.0", "*", ""):
+            actual_ip = get_ip()
+            # If worker is on the same machine, use loopback for reliability
+            # This handles cases where external IP (e.g., VPN) doesn't support
+            # local loopback connections
+            if worker_ip is not None and worker_ip == actual_ip:
+                logger.debug(
+                    "Worker IP %s matches controller IP, using 127.0.0.1",
+                    worker_ip,
+                )
+                return f"127.0.0.1:{port}"
+            return f"{actual_ip}:{port}"
+        return bind_url
+
+    def _check_socket_has_pending(self, socket) -> int:
+        """Check if socket has pending messages.
+
+        Returns:
+            1 if socket has pending messages, 0 otherwise
+        """
+        try:
+            events = socket.get(zmq.EVENTS)  # type: ignore[attr-defined]
+            # Check if POLLIN flag is set (indicates readable/pending messages)
+            has_pending = 1 if (events & zmq.POLLIN) else 0  # type: ignore[attr-defined]
+            return has_pending
+        except Exception as e:
+            logger.error(f"Error checking socket pending status: {e}")
+            return 0
+
+    async def handle_batched_push_request(self, socket) -> Optional[MsgBase]:
+        while True:
+            parts = await socket.recv_multipart()
+            part_count = len(parts)
+            with SocketMetricsContext(self, SocketType.PULL, part_count):
+                for part in parts:
+                    # Parse message based on format
+                    if part.startswith(b"{"):
+                        # JSON format - typically from external systems
+                        # like Mooncake
+                        msg_dict = json.loads(part)
+                        msg = msgspec.convert(msg_dict, type=Msg)
+                    else:
+                        # MessagePack format - internal LMCache communication
+                        msg = msgspec.msgpack.decode(part, type=Msg)
+                    if isinstance(msg, WorkerMsg):
+                        await self.handle_worker_message(msg)
+
+                    # FIXME(Jiayi): The abstraction of control messages
+                    # might not be necessary.
+                    # elif isinstance(msg, ControlMsg):
+                    #    await self.issue_control_message(msg)
+                    elif isinstance(msg, OrchMsg):
+                        await self.handle_orchestration_message(msg)
+                    else:
+                        logger.error(f"Unknown message type: {type(msg)}")
+
+    async def handle_batched_req_request(self, socket) -> Optional[MsgBase]:
+        """Handle requests on ROUTER socket.
+
+        ROUTER socket receives multi-part messages:
+        [identity, empty_frame, payload]
+        and must reply with the same identity frame.
+        """
+        while True:
+            frames = await socket.recv_multipart()
+            with SocketMetricsContext(self, SocketType.REPLY):
+                identity = None
+                try:
+                    # ROUTER socket: [identity, empty_frame, payload]
+                    if len(frames) < 3:
+                        logger.error(
+                            "Invalid ROUTER message format, expected >= 3 frames, "
+                            "got %d",
+                            len(frames),
+                        )
+                        continue
+                    identity = frames[0]
+                    # frames[1] is empty delimiter
+                    part = frames[2]
+
+                    # Parse message based on format
+                    if part.startswith(b"{"):
+                        # JSON format - typically from external systems like Mooncake
+                        msg_dict = json.loads(part)
+                        msg = msgspec.convert(msg_dict, type=Msg)
+                    else:
+                        # MessagePack format - internal LMCache communication
+                        msg = msgspec.msgpack.decode(part, type=Msg)
+
+                    if isinstance(msg, WorkerReqMsg):
+                        ret_msg = await self.handle_worker_req_message(msg)
+                        # Reply with identity frame for ROUTER socket
+                        await socket.send_multipart(
+                            [identity, b"", msgspec.msgpack.encode(ret_msg)]
+                        )
+                    else:
+                        logger.error("Unknown message type: %s", type(msg))
+                        err_msg = ErrorMsg(error=f"Unknown message type: {type(msg)}")
+                        await socket.send_multipart(
+                            [identity, b"", msgspec.msgpack.encode(err_msg)]
+                        )
+                except (
+                    json.JSONDecodeError,
+                    msgspec.DecodeError,
+                    msgspec.ValidationError,
+                    zmq.ZMQError,
+                ) as e:
+                    logger.error("Error handling request message: %s", e, exc_info=True)
+                    err_msg = ErrorMsg(error=str(e))
+                    # Try to reply with error if we have identity
+                    if identity is not None:
+                        await socket.send_multipart(
+                            [identity, b"", msgspec.msgpack.encode(err_msg)]
+                        )
+
+    async def handle_heartbeat_request(self, socket) -> None:
+        """Handle heartbeat requests on dedicated ROUTER socket.
+
+        This runs on a separate socket to ensure heartbeats are processed
+        without being blocked by other requests.
+
+        ROUTER socket receives multi-part messages:
+        [identity, empty_frame, payload]
+        """
+        logger.info("Heartbeat handler task started, waiting for heartbeat requests...")
+        while True:
+            frames = await socket.recv_multipart()
+            logger.debug("Received heartbeat request with %d frames", len(frames))
+            with SocketMetricsContext(self, SocketType.REPLY):
+                identity = None
+                try:
+                    # ROUTER socket: [identity, empty_frame, payload]
+                    if len(frames) < 3:
+                        logger.error(
+                            "Invalid heartbeat ROUTER message format, "
+                            "expected >= 3 frames, got %d",
+                            len(frames),
+                        )
+                        continue
+                    identity = frames[0]
+                    part = frames[2]
+
+                    if part.startswith(b"{"):
+                        msg_dict = json.loads(part)
+                        msg = msgspec.convert(msg_dict, type=Msg)
+                    else:
+                        msg = msgspec.msgpack.decode(part, type=Msg)
+
+                    if isinstance(msg, HeartbeatMsg):
+                        ret_msg = await self.reg_controller.heartbeat(msg)
+                        await socket.send_multipart(
+                            [identity, b"", msgspec.msgpack.encode(ret_msg)]
+                        )
+                    else:
+                        logger.error(
+                            "Unexpected message type on heartbeat socket: %s",
+                            type(msg),
+                        )
+                        err_msg = ErrorMsg(
+                            error=f"Expected HeartbeatMsg, got {type(msg)}"
+                        )
+                        await socket.send_multipart(
+                            [identity, b"", msgspec.msgpack.encode(err_msg)]
+                        )
+                except (
+                    json.JSONDecodeError,
+                    msgspec.DecodeError,
+                    msgspec.ValidationError,
+                    zmq.ZMQError,
+                ) as e:
+                    logger.error(
+                        "Error handling heartbeat request: %s", e, exc_info=True
+                    )
+                    err_msg = ErrorMsg(error=str(e))
+                    # Try to reply with error if we have identity
+                    if identity is not None:
+                        await socket.send_multipart(
+                            [identity, b"", msgspec.msgpack.encode(err_msg)]
+                        )
+
+    async def health_check(self):
+        while True:
+            await asyncio.sleep(self.health_check_interval)
+            worker_infos = self.reg_controller.registry.get_all_worker_infos_cached(1)
+            for worker_info in worker_infos:
+                if (
+                    time.time() - worker_info.last_heartbeat_time
+                    > self.lmcache_worker_timeout
+                ):
+                    logger.warning(
+                        "Worker %s_%s last heartbeat time: %s, "
+                        "current time: %s, more than %s seconds",
+                        worker_info.instance_id,
+                        worker_info.worker_id,
+                        worker_info.last_heartbeat_time,
+                        time.time(),
+                        self.lmcache_worker_timeout,
+                    )
+                    # Perform a full deregister to clean up all associated resources.
+                    deregister_msg = DeRegisterMsg(
+                        instance_id=worker_info.instance_id,
+                        worker_id=worker_info.worker_id,
+                        ip=worker_info.ip,
+                        port=worker_info.port,
+                    )
+                    await self.reg_controller.deregister(deregister_msg)
+
+    async def start_all(self):
+        tasks = []
+        if self.controller_urls["reply"] is not None:
+            logger.info(
+                "Starting reply request handler on %s", self.controller_urls["reply"]
+            )
+            tasks.append(self.handle_batched_req_request(self.controller_reply_socket))
+        if self.controller_heartbeat_socket is not None:
+            logger.info(
+                "Starting heartbeat request handler on %s",
+                self.controller_urls.get("heartbeat"),
+            )
+            tasks.append(
+                self.handle_heartbeat_request(self.controller_heartbeat_socket)
+            )
+        else:
+            logger.warning("Heartbeat socket is None, heartbeat handler not started!")
+        logger.info("Starting pull request handler on %s", self.controller_urls["pull"])
+        tasks.append(self.handle_batched_push_request(self.controller_pull_socket))
+        await asyncio.gather(
+            *tasks,
+            return_exceptions=True,
+        )
+
+    def close(self):
+        """Clean up all resources owned by the controller manager."""
+        if hasattr(self, "controller_pull_socket"):
+            self.controller_pull_socket.close()
+        if hasattr(self, "controller_reply_socket"):
+            self.controller_reply_socket.close()
+        if hasattr(self, "zmq_context"):
+            self.zmq_context.destroy()

lmcache/v1/cache_controller/controllers/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.cache_controller.controllers.kv_controller import KVController
+from lmcache.v1.cache_controller.controllers.registration_controller import (  # noqa: E501
+    RegistrationController,
+)
+
+__all__ = [
+    "KVController",
+    "RegistrationController",
+]