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

lmcache/v1/cache_controller/executor.py

@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Union
+import asyncio
+import uuid
+
+# Third Party
+import msgspec
+import zmq.asyncio
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.cache_controller.message import (  # noqa: E501
+    CheckFinishMsg,
+    CheckFinishRetMsg,
+    ClearMsg,
+    ClearRetMsg,
+    ClearWorkerMsg,
+    CompressMsg,
+    CompressRetMsg,
+    CompressWorkerMsg,
+    DecompressMsg,
+    DecompressRetMsg,
+    DecompressWorkerMsg,
+    ErrorMsg,
+    HealthMsg,
+    HealthRetMsg,
+    HealthWorkerMsg,
+    HealthWorkerRetMsg,
+    MoveMsg,
+    MoveRetMsg,
+    MoveWorkerMsg,
+    Msg,
+    MsgBase,
+    PinMsg,
+    PinRetMsg,
+    PinWorkerMsg,
+)
+
+logger = init_logger(__name__)
+
+
+# NOTE (Jiayi): `LMCacheClusterExecutor` might need to be in different processes
+# in the future for the sake of performance.
+# NOTE (Jiayi): Also, consider scaling up the number of cluster executors
+# in the future.
+# TODO (Jiayi): need better error handling
+class LMCacheClusterExecutor:
+    """
+    LMCache Cluster Executor class to handle the execution of cache operations.
+    """
+
+    def __init__(self, reg_controller):
+        """
+        Initialize the LMCache Executor with a cache instance.
+
+        :param lmcache_instance_id: lmcache_instance_id
+        """
+        self.reg_controller = reg_controller
+
+    async def clear(self, msg: ClearMsg) -> Union[ClearRetMsg, ErrorMsg]:
+        """
+        Execute a clear cache operation with error handling.
+        """
+        instance_id = msg.instance_id
+        location = msg.location
+
+        worker_ids = self.reg_controller.get_workers(instance_id)
+        assert worker_ids is not None
+        sockets = []
+        serialized_msgs = []
+        for worker_id in worker_ids:
+            socket = self.reg_controller.get_socket(instance_id, worker_id)
+            if socket is None:
+                return ErrorMsg(
+                    error=(
+                        f"Worker {worker_id} not registered for instance {instance_id}"
+                    )
+                )
+            sockets.append(socket)
+
+            # TODO(Jiayi): Need a way to trak event_id -> worker_event_id mapping
+            # Also, we need to track worker_event_id status
+            worker_event_id = f"Worker{worker_id}{msg.event_id}"
+            serialized_msg = msgspec.msgpack.encode(
+                ClearWorkerMsg(
+                    worker_event_id=worker_event_id,
+                    location=location,
+                )
+            )
+            serialized_msgs.append(serialized_msg)
+        serialized_results = await self.execute_workers(
+            sockets=sockets,
+            serialized_msgs=serialized_msgs,
+        )
+
+        num_tokens_list = []
+        for i, serialized_result in enumerate(serialized_results):
+            result = msgspec.msgpack.decode(serialized_result, type=Msg)
+            num_tokens_list.append(result.num_tokens)
+
+        # TODO(Jiayi): Need to ensure cache consistency across workers.
+        assert len(set(num_tokens_list)) == 1, (
+            "The number of tokens cleared should be the same across all workers."
+        )
+
+        return ClearRetMsg(event_id=msg.event_id, num_tokens=num_tokens_list[0])
+
+    async def pin(self, msg: PinMsg) -> Union[PinRetMsg, ErrorMsg]:
+        """
+        Execute a pin cache operation with error handling.
+        """
+        instance_id = msg.instance_id
+        tokens = msg.tokens
+        location = msg.location
+
+        worker_ids = self.reg_controller.get_workers(instance_id)
+        assert worker_ids is not None
+        sockets = []
+        serialized_msgs = []
+        for worker_id in worker_ids:
+            socket = self.reg_controller.get_socket(instance_id, worker_id)
+            if socket is None:
+                return ErrorMsg(
+                    error=(
+                        f"Worker {worker_id} not registered for instance {instance_id}"
+                    )
+                )
+            sockets.append(socket)
+
+            # TODO(Jiayi): Need a way to trak event_id -> worker_event_id mapping
+            # Also, we need to track worker_event_id status
+            worker_event_id = f"Worker{worker_id}{msg.event_id}"
+            serialized_msg = msgspec.msgpack.encode(
+                PinWorkerMsg(
+                    worker_event_id=worker_event_id,
+                    tokens=tokens,
+                    location=location,
+                )
+            )
+            serialized_msgs.append(serialized_msg)
+        serialized_results = await self.execute_workers(
+            sockets=sockets,
+            serialized_msgs=serialized_msgs,
+        )
+
+        num_tokens_list = []
+        for i, serialized_result in enumerate(serialized_results):
+            result = msgspec.msgpack.decode(serialized_result, type=Msg)
+            num_tokens_list.append(result.num_tokens)
+
+        # TODO(Jiayi): Need to ensure cache consistency across workers.
+        assert len(set(num_tokens_list)) == 1, (
+            "The number of tokens pinned should be the same across all workers."
+        )
+
+        return PinRetMsg(event_id=msg.event_id, num_tokens=num_tokens_list[0])
+
+    async def compress(self, msg: CompressMsg) -> Union[CompressRetMsg, ErrorMsg]:
+        """
+        Execute a compress operation with error handling.
+        """
+        event_id = msg.event_id
+        instance_id = msg.instance_id
+        method = msg.method
+        location = msg.location
+        tokens = msg.tokens
+
+        worker_ids = self.reg_controller.get_workers(instance_id)
+        assert worker_ids is not None
+
+        # TODO(Jiayi): Currently, we do not support PP or heterogeneous TP.
+        # NOTE(Jiayi): The TP ranks are already sorted in registration_controller.
+
+        sockets = []
+        serialized_msgs = []
+        for worker_id in worker_ids:
+            socket = self.reg_controller.get_socket(instance_id, worker_id)
+
+            if socket is None:
+                return ErrorMsg(
+                    error=(
+                        f"Worker {worker_id} not registered for "
+                        f"instance {instance_id} or "
+                    )
+                )
+            sockets.append(socket)
+
+            worker_event_id = f"CompressWorker{worker_id}{str(uuid.uuid4())}"
+            serialized_msg = msgspec.msgpack.encode(
+                CompressWorkerMsg(
+                    worker_event_id=worker_event_id,
+                    method=method,
+                    location=location,
+                    tokens=tokens,
+                )
+            )
+            serialized_msgs.append(serialized_msg)
+            logger.debug(
+                f"Sending compress operation to worker ({instance_id}, {worker_id})"
+            )
+        serialized_results = await self.execute_workers(
+            sockets=sockets,
+            serialized_msgs=serialized_msgs,
+        )
+
+        num_tokens_list = []
+        for serialized_result in serialized_results:
+            result = msgspec.msgpack.decode(serialized_result, type=Msg)
+            num_tokens_list.append(result.num_tokens)
+
+        # TODO(Jiayi): Need to ensure cache consistency across workers.
+        assert len(set(num_tokens_list)) == 1, (
+            "The number of tokens compressed should be the same across all workers."
+        )
+
+        return CompressRetMsg(
+            event_id=event_id,
+            num_tokens=num_tokens_list[0],
+        )
+
+    async def decompress(self, msg: DecompressMsg) -> Union[DecompressRetMsg, ErrorMsg]:
+        """
+        Execute a decompress operation with error handling.
+        """
+        event_id = msg.event_id
+        instance_id = msg.instance_id
+        method = msg.method
+        location = msg.location
+        tokens = msg.tokens
+
+        worker_ids = self.reg_controller.get_workers(instance_id)
+        assert worker_ids is not None
+
+        sockets = []
+        serialized_msgs = []
+        for worker_id in worker_ids:
+            socket = self.reg_controller.get_socket(instance_id, worker_id)
+
+            if socket is None:
+                return ErrorMsg(
+                    error=(
+                        f"Worker {worker_id} not registered for "
+                        f"instance {instance_id} or "
+                    )
+                )
+            sockets.append(socket)
+
+            worker_event_id = f"DecompressWorker{worker_id}{str(uuid.uuid4())}"
+            serialized_msg = msgspec.msgpack.encode(
+                DecompressWorkerMsg(
+                    worker_event_id=worker_event_id,
+                    method=method,
+                    location=location,
+                    tokens=tokens,
+                )
+            )
+            serialized_msgs.append(serialized_msg)
+            logger.debug(
+                f"Sending decompress operation to worker ({instance_id}, {worker_id})"
+            )
+        serialized_results = await self.execute_workers(
+            sockets=sockets,
+            serialized_msgs=serialized_msgs,
+        )
+
+        num_tokens_list = []
+        for serialized_result in serialized_results:
+            result = msgspec.msgpack.decode(serialized_result, type=Msg)
+            num_tokens_list.append(result.num_tokens)
+
+        assert len(set(num_tokens_list)) == 1, (
+            "The number of tokens decompressed should be the same across all workers."
+        )
+
+        return DecompressRetMsg(
+            event_id=event_id,
+            num_tokens=num_tokens_list[0],
+        )
+
+    async def move(self, msg: MoveMsg) -> Union[MoveRetMsg, ErrorMsg]:
+        """
+        Execute a move cache operation with error handling.
+        """
+        # NOTE(Jiayi): Currently we assume the transfer is push-based.
+        src_instance_id = msg.old_position[0]
+        dst_instance_id = msg.new_position[0]
+
+        src_worker_ids = self.reg_controller.get_workers(src_instance_id)
+        assert src_worker_ids is not None
+        dst_worker_ids = self.reg_controller.get_workers(dst_instance_id)
+        assert dst_worker_ids is not None
+
+        # TODO(Jiayi): Currently, we do not support PP or heterogeneous TP.
+        # NOTE(Jiayi): The TP ranks are already sorted in registration_controller.
+
+        sockets = []
+        serialized_msgs = []
+        for src_worker_id, dst_worker_id in zip(
+            src_worker_ids, dst_worker_ids, strict=False
+        ):
+            socket = self.reg_controller.get_socket(src_instance_id, src_worker_id)
+            dst_url = self.reg_controller.get_peer_init_url(
+                dst_instance_id, dst_worker_id
+            )
+
+            if socket is None or dst_url is None:
+                return ErrorMsg(
+                    error=(
+                        f"Src worker {src_worker_id} not registered for "
+                        f"instance {src_instance_id} or "
+                        f"dst worker {dst_worker_id} not registered for "
+                        f"instance {dst_instance_id} or P2P is not enabled."
+                    )
+                )
+            sockets.append(socket)
+
+            worker_event_id = f"MoveWorker{src_worker_id}{str(uuid.uuid4())}"
+            serialized_msg = msgspec.msgpack.encode(
+                MoveWorkerMsg(
+                    worker_event_id=worker_event_id,
+                    old_position=msg.old_position[1],
+                    new_position=(dst_url, msg.new_position[1]),
+                    tokens=msg.tokens,
+                    copy=msg.copy,
+                )
+            )
+            serialized_msgs.append(serialized_msg)
+            logger.debug(
+                f"Sending move operation to worker ({src_instance_id}, {src_worker_id})"
+            )
+        serialized_results = await self.execute_workers(
+            sockets=sockets,
+            serialized_msgs=serialized_msgs,
+        )
+
+        num_tokens_list = []
+        for serialized_result in serialized_results:
+            result = msgspec.msgpack.decode(serialized_result, type=Msg)
+            num_tokens_list.append(result.num_tokens)
+
+        # TODO(Jiayi): Need to ensure cache consistency across workers.
+        assert len(set(num_tokens_list)) == 1, (
+            "The number of tokens moved should be the same across all workers."
+        )
+
+        return MoveRetMsg(
+            event_id=msg.event_id,
+            num_tokens=num_tokens_list[0],
+        )
+
+    async def health(self, msg: HealthMsg) -> Union[HealthRetMsg, ErrorMsg]:
+        """
+        Execute a compress operation with error handling.
+        """
+        instance_id = msg.instance_id
+
+        worker_ids = self.reg_controller.get_workers(instance_id)
+        if worker_ids is None:
+            return ErrorMsg(error=f"No workers found for instance {instance_id}")
+
+        # TODO(Jiayi): Currently, we do not support PP or heterogeneous TP.
+        # NOTE(Jiayi): The TP ranks are already sorted in registration_controller.
+
+        sockets = []
+        serialized_msgs = []
+        for worker_id in worker_ids:
+            socket = self.reg_controller.get_socket(instance_id, worker_id)
+
+            if socket is None:
+                return ErrorMsg(
+                    error=(
+                        f"Worker {worker_id} not registered for "
+                        f"instance {instance_id} or socket not found"
+                    )
+                )
+            sockets.append(socket)
+
+            worker_event_id = f"HealthWorker{worker_id}{str(uuid.uuid4())}"
+            serialized_msg = msgspec.msgpack.encode(
+                HealthWorkerMsg(
+                    worker_event_id=worker_event_id,
+                )
+            )
+            serialized_msgs.append(serialized_msg)
+            logger.debug(
+                f"Sending health check operation to worker ({instance_id}, {worker_id})"
+            )
+
+        # Collect results from all workers
+        serialized_results = await self.execute_workers(
+            sockets=sockets,
+            serialized_msgs=serialized_msgs,
+        )
+
+        # Process results
+        error_codes = {}
+        for i, serialized_result in enumerate(serialized_results):
+            try:
+                result = msgspec.msgpack.decode(serialized_result, type=Msg)
+                if isinstance(result, HealthWorkerRetMsg):
+                    error_codes[worker_ids[i]] = result.error_code
+                elif isinstance(result, ErrorMsg):
+                    error_codes[worker_ids[i]] = -1001  # Worker returned error
+                else:
+                    error_codes[worker_ids[i]] = -1002  # Unexpected response
+            except Exception as e:
+                logger.error(
+                    f"Failed to parse health response from worker "
+                    f"{worker_ids[i]}: {str(e)}"
+                )
+                error_codes[worker_ids[i]] = -1003  # Failed to parse response
+
+        return HealthRetMsg(
+            event_id=msg.event_id,
+            error_codes=error_codes,
+        )
+
+    async def check_finish(
+        self, msg: CheckFinishMsg
+    ) -> Union[CheckFinishRetMsg, ErrorMsg]:
+        raise NotImplementedError
+
+    # TODO(Jiayi): need to make the types more specific
+    async def execute(self, operation: str, msg: MsgBase) -> MsgBase:
+        """
+        Execute a cache operation with error handling.
+
+        :param operation: The operation to execute
+        (e.g., 'clear').
+        :param msg: The message containing the operation details.
+        :return: The result of the operation or an error message.
+        """
+        try:
+            method = getattr(self, operation)
+            return await method(msg)
+        except AttributeError:
+            return ErrorMsg(error=f"Operation '{operation}' is not supported.")
+        except Exception as e:
+            return ErrorMsg(error=str(e))
+
+    async def execute_workers(
+        self,
+        sockets: list[zmq.asyncio.Socket],
+        serialized_msgs: list[bytes],
+    ) -> list[bytes]:
+        """
+        Execute a list of serialized messages on the given sockets.
+        :param sockets: The list of sockets to send the messages to.
+        :param serialized_msgs: The list of serialized messages to send.
+        :return: A list of serialized results received from the sockets.
+        """
+        tasks = []
+        for socket, serialized_msg in zip(sockets, serialized_msgs, strict=False):
+
+            async def send_and_receive(s, msg):
+                await s.send(msg)
+                return await s.recv()
+
+            tasks.append(send_and_receive(socket, serialized_msg))
+
+        serialized_results = await asyncio.gather(*tasks)
+        return serialized_results

lmcache/v1/cache_controller/frontend/static/css/style.css

@@ -0,0 +1,201 @@
+/* Controller Dashboard Styles */
+body {
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif;
+    background-color: #f8f9fa;
+}
+
+.container-fluid {
+    max-width: 1400px;
+    margin: 0 auto;
+}
+
+header {
+    border-radius: 0.5rem;
+    margin-bottom: 1rem;
+}
+
+.nav-tabs {
+    border-bottom: 2px solid #dee2e6;
+}
+
+.nav-tabs .nav-link {
+    border: none;
+    color: #495057;
+    font-weight: 500;
+    padding: 0.75rem 1.5rem;
+}
+
+.nav-tabs .nav-link:hover {
+    border: none;
+    color: #0d6efd;
+}
+
+.nav-tabs .nav-link.active {
+    color: #0d6efd;
+    border-bottom: 3px solid #0d6efd;
+    background-color: transparent;
+}
+
+.tab-content {
+    background-color: white;
+    border-radius: 0 0 0.5rem 0.5rem;
+    box-shadow: 0 0.125rem 0.25rem rgba(0, 0, 0, 0.075);
+}
+
+.card {
+    border: 1px solid rgba(0, 0, 0, 0.125);
+    border-radius: 0.5rem;
+    box-shadow: 0 0.125rem 0.25rem rgba(0, 0, 0, 0.075);
+}
+
+.card-title {
+    color: #495057;
+    font-weight: 600;
+    margin-bottom: 1rem;
+}
+
+.table-responsive {
+    margin-top: 1rem;
+}
+
+.table {
+    margin-bottom: 0;
+}
+
+.table thead th {
+    background-color: #f8f9fa;
+    border-bottom: 2px solid #dee2e6;
+    font-weight: 600;
+    color: #495057;
+}
+
+.table tbody tr:hover {
+    background-color: #f8f9fa;
+}
+
+.btn {
+    border-radius: 0.375rem;
+    font-weight: 500;
+}
+
+.btn-primary {
+    background-color: #0d6efd;
+    border-color: #0d6efd;
+}
+
+.btn-primary:hover {
+    background-color: #0b5ed7;
+    border-color: #0a58ca;
+}
+
+.btn-outline-primary {
+    color: #0d6efd;
+    border-color: #0d6efd;
+}
+
+.btn-outline-primary:hover {
+    background-color: #0d6efd;
+    border-color: #0d6efd;
+}
+
+.form-control, .form-select {
+    border-radius: 0.375rem;
+    border: 1px solid #ced4da;
+}
+
+.form-control:focus, .form-select:focus {
+    border-color: #86b7fe;
+    box-shadow: 0 0 0 0.25rem rgba(13, 110, 253, 0.25);
+}
+
+pre {
+    background-color: #f8f9fa;
+    border: 1px solid #dee2e6;
+    border-radius: 0.375rem;
+    padding: 1rem;
+    white-space: pre-wrap;
+    word-wrap: break-word;
+    font-family: 'Courier New', Courier, monospace;
+    font-size: 0.875rem;
+}
+
+.badge {
+    font-weight: 500;
+    padding: 0.35em 0.65em;
+}
+
+.badge.bg-success {
+    background-color: #198754 !important;
+}
+
+.badge.bg-warning {
+    background-color: #ffc107 !important;
+}
+
+.badge.bg-danger {
+    background-color: #dc3545 !important;
+}
+
+.badge.bg-info {
+    background-color: #0dcaf0 !important;
+}
+
+.spinner-border {
+    margin: 1rem auto;
+    display: block;
+}
+
+.modal-content {
+    border-radius: 0.5rem;
+}
+
+.modal-header {
+    border-bottom: 1px solid #dee2e6;
+}
+
+.modal-footer {
+    border-top: 1px solid #dee2e6;
+}
+
+/* Status indicators */
+.status-active {
+    color: #198754;
+}
+
+.status-inactive {
+    color: #6c757d;
+}
+
+.status-error {
+    color: #dc3545;
+}
+
+.status-warning {
+    color: #ffc107;
+}
+
+/* Responsive adjustments */
+@media (max-width: 768px) {
+    .container-fluid {
+        padding: 0.5rem;
+    }
+    
+    .nav-tabs .nav-link {
+        padding: 0.5rem 0.75rem;
+        font-size: 0.875rem;
+    }
+    
+    .tab-content {
+        padding: 1rem;
+    }
+}
+
+/* Animation for refreshing */
+@keyframes spin {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+}
+
+.refreshing {
+    animation: spin 1s linear infinite;
+}