# 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()