lmcache--lmcache/lmcache/v1/cache_controller/controllers/registration_controller.py

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