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

lmcache/v1/cache_controller/controllers/registration_controller.py

@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional
+import time
+
+# Third Party
+import zmq
+import zmq.asyncio
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.commands import FullSyncCommand, HeartbeatCommand
+from lmcache.v1.cache_controller.message import (
+    DeRegisterMsg,
+    HealthMsg,
+    HealthRetMsg,
+    HeartbeatMsg,
+    HeartbeatRetMsg,
+    QueryInstMsg,
+    QueryInstRetMsg,
+    QueryWorkerInfoMsg,
+    QueryWorkerInfoRetMsg,
+    RegisterMsg,
+    RegisterRetMsg,
+)
+from lmcache.v1.cache_controller.observability import PrometheusLogger
+from lmcache.v1.cache_controller.utils import RegistryTree
+from lmcache.v1.rpc_utils import (
+    close_zmq_socket,
+    get_zmq_context,
+    get_zmq_socket,
+)
+
+logger = init_logger(__name__)
+
+
+class RegistrationController:
+    def __init__(self):
+        # Central registry tree managing all instances and workers
+        self.registry = RegistryTree()
+        self._setup_metrics()
+
+    def _setup_metrics(self):
+        prometheus_logger = PrometheusLogger.GetInstanceOrNone()
+        if prometheus_logger is not None:
+            prometheus_logger.registered_workers_count.set_function(
+                lambda: len(self.registry.get_all_worker_infos_cached())
+            )
+
+    def post_init(self, kv_controller, cluster_executor):
+        """
+        Post initialization of the Registration Controller.
+        """
+        self.kv_controller = kv_controller
+        self.cluster_executor = cluster_executor
+
+    def get_socket(
+        self, instance_id: str, worker_id: int
+    ) -> Optional[zmq.asyncio.Socket]:
+        """
+        Get the socket for a given instance and worker ID.
+        """
+        worker_node = self.registry.get_worker(instance_id, worker_id)
+        if worker_node is None:
+            logger.warning(
+                "Instance-worker %s not registered", (instance_id, worker_id)
+            )
+            return None
+        return worker_node.socket
+
+    def get_peer_init_url(self, instance_id: str, worker_id: int) -> Optional[str]:
+        """
+        Get the URL for a given instance and worker ID.
+        """
+        worker_node = self.registry.get_worker(instance_id, worker_id)
+        if worker_node is None:
+            logger.warning(
+                "Instance-worker %s not registered or P2P is not used",
+                (instance_id, worker_id),
+            )
+            return None
+        return worker_node.peer_init_url
+
+    def get_workers(self, instance_id: str) -> list[int]:
+        """
+        Get worker ids given an instance id.
+        """
+        return self.registry.get_worker_ids(instance_id)
+
+    async def get_instance_id(self, msg: QueryInstMsg) -> QueryInstRetMsg:
+        """
+        Get the instance id given an ip address.
+        """
+        ip = msg.ip
+        event_id = msg.event_id
+        instance_node = self.registry.get_instance_by_ip(ip)
+        if instance_node is None:
+            logger.warning("Instance not registered for IP %s", ip)
+            return QueryInstRetMsg(instance_id=None, event_id=event_id)
+        return QueryInstRetMsg(instance_id=instance_node.instance_id, event_id=event_id)
+
+    async def register(
+        self, msg: RegisterMsg, extra_config: Optional[dict[str, str]] = None
+    ) -> RegisterRetMsg:
+        """
+        Register a new instance-worker connection mapping.
+
+        Args:
+            msg: RegisterMsg from worker
+            extra_config: Optional extra configuration to return to worker,
+                          e.g., {"heartbeat_url": "tcp://...:8082"}
+
+        Returns:
+            RegisterRetMsg with extra_config for worker initialization
+        """
+        instance_id = msg.instance_id
+        worker_id = msg.worker_id
+        ip = msg.ip
+        port = msg.port
+        url = f"{ip}:{port}"
+
+        # prevent duplicate registration
+        existing_worker = self.registry.get_worker(instance_id, worker_id)
+        if existing_worker is not None:
+            logger.warning(
+                "Instance-worker %s already registered, skip registration",
+                (instance_id, worker_id),
+            )
+            self.registry.clear_worker_kv(instance_id, worker_id)
+            return (
+                RegisterRetMsg()
+                if extra_config is None
+                else RegisterRetMsg(extra_config=extra_config)
+            )
+
+        peer_init_url = msg.peer_init_url
+        if peer_init_url is None:
+            logger.info(
+                "peer init url of %s is None, only register when p2p is used.",
+                (instance_id, worker_id),
+            )
+
+        context = get_zmq_context()
+        socket = get_zmq_socket(
+            context,
+            url,
+            protocol="tcp",
+            role=zmq.REQ,  # type: ignore[attr-defined]
+            bind_or_connect="connect",
+        )
+
+        # Register worker in the tree
+        self.registry.register_worker(
+            instance_id=instance_id,
+            worker_id=worker_id,
+            ip=ip,
+            port=port,
+            peer_init_url=peer_init_url,
+            socket=socket,
+            registration_time=time.time(),
+        )
+
+        logger.info(
+            "Registered instance-worker %s with URL %s", (instance_id, worker_id), url
+        )
+        return (
+            RegisterRetMsg()
+            if extra_config is None
+            else RegisterRetMsg(extra_config=extra_config)
+        )
+
+    async def deregister(self, msg: DeRegisterMsg) -> None:
+        """
+        Deregister an instance-worker connection mapping.
+        """
+        instance_id = msg.instance_id
+        worker_id = msg.worker_id
+
+        worker_node = self.registry.deregister_worker(instance_id, worker_id)
+        if worker_node is None:
+            logger.warning(
+                "Instance-worker %s not registered", (instance_id, worker_id)
+            )
+            return
+
+        # Close socket
+        if worker_node.socket is not None:
+            close_zmq_socket(worker_node.socket)
+
+        logger.info("Deregistered instance-worker %s", (instance_id, worker_id))
+
+    async def health(self, msg: HealthMsg) -> HealthRetMsg:
+        """
+        Check the health of the lmcache worker.
+        """
+        return await self.cluster_executor.execute(
+            "health",
+            msg,
+        )
+
+    async def heartbeat(self, msg: HeartbeatMsg) -> HeartbeatRetMsg:
+        """
+        Heartbeat from lmcache worker (REQ-REP mode).
+
+        Returns HeartbeatRetMsg with optional commands for the worker to execute.
+        Commands are executed sequentially by the worker.
+        """
+        instance_id = msg.instance_id
+        worker_id = msg.worker_id
+        success = self.registry.update_heartbeat(instance_id, worker_id, time.time())
+
+        commands: list[HeartbeatCommand] = []
+
+        if not success:
+            logger.warning(
+                "%s has not been registered, re-register the worker.",
+                (instance_id, worker_id),
+            )
+            # re-register the worker
+            register_msg = RegisterMsg(
+                instance_id=msg.instance_id,
+                worker_id=msg.worker_id,
+                ip=msg.ip,
+                port=msg.port,
+                peer_init_url=msg.peer_init_url,
+            )
+            await self.register(register_msg)
+            # New worker needs full sync
+            commands.append(FullSyncCommand(reason="worker_re_registered"))
+        else:
+            # Check if full sync is needed (e.g., controller restart)
+            if self.kv_controller is not None and hasattr(
+                self.kv_controller, "full_sync_tracker"
+            ):
+                should_sync, reason = (
+                    self.kv_controller.full_sync_tracker.should_request_full_sync(
+                        instance_id, worker_id
+                    )
+                )
+                if should_sync:
+                    commands.append(FullSyncCommand(reason=reason))
+
+        return HeartbeatRetMsg(commands=commands)
+
+    async def query_worker_info(self, msg: QueryWorkerInfoMsg) -> QueryWorkerInfoRetMsg:
+        """
+        Query worker info.
+        """
+        event_id = msg.event_id
+        worker_infos = []
+
+        # Handle special case: instance_id = "all"
+        if msg.instance_id == "all":
+            # Get all worker infos from the registry
+            worker_infos = self.registry.get_all_worker_infos_cached()
+            # If specific worker_ids are requested, filter the results
+            if msg.worker_ids is not None and len(msg.worker_ids) > 0:
+                worker_infos = [
+                    worker_info
+                    for worker_info in worker_infos
+                    if worker_info.worker_id in msg.worker_ids
+                ]
+            return QueryWorkerInfoRetMsg(event_id=event_id, worker_infos=worker_infos)
+
+        # Normal case: query specific instance
+        instance_node = self.registry.get_instance(msg.instance_id)
+        if instance_node is None:
+            logger.warning("instance %s not registered.", msg.instance_id)
+        else:
+            worker_ids = msg.worker_ids
+            if worker_ids is None or len(worker_ids) == 0:
+                worker_ids = instance_node.get_worker_ids()
+            for worker_id in worker_ids:
+                worker_node = instance_node.get_worker(worker_id)
+                if worker_node is not None:
+                    worker_infos.append(worker_node.to_worker_info(msg.instance_id))
+                else:
+                    logger.warning(
+                        "worker %s not registered.", (msg.instance_id, worker_id)
+                    )
+
+        return QueryWorkerInfoRetMsg(event_id=event_id, worker_infos=worker_infos)