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

lmcache/tools/controller_benchmark/constants.py

@@ -0,0 +1,10 @@
+"""Constants for LMCache Controller ZMQ Benchmark"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# ZMQ socket configuration
+DEFAULT_SEND_TIMEOUT_MS = 5000
+DEFAULT_RECV_TIMEOUT_MS = 5000
+DEFAULT_SEND_HWM = 10000
+DEFAULT_BATCH_SEND_SIZE = 50
+DEFAULT_OP_DISTRIBUTION_BASE = 1000

lmcache/tools/controller_benchmark/handlers/__init__.py

@@ -0,0 +1,46 @@
+"""Operation handlers with dynamic discovery and registration"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import Dict
+import importlib
+import inspect
+import pkgutil
+
+# Local
+from .base import OperationHandler
+
+# Operation handler registry
+OPERATION_HANDLERS: Dict[str, OperationHandler] = {}
+
+
+def _discover_and_register_handlers():
+    """Dynamically discover and register all operation handlers"""
+    # Get the current package
+    package = __package__
+    package_path = __path__
+
+    # Iterate through all modules in the handlers package
+    for _, module_name, _ in pkgutil.iter_modules(package_path):
+        # Skip base module
+        if module_name == "base":
+            continue
+
+        # Import the module
+        module = importlib.import_module("." + module_name, package=package)
+
+        # Find all classes that inherit from OperationHandler
+        for name, obj in inspect.getmembers(module, inspect.isclass):
+            if (
+                issubclass(obj, OperationHandler)
+                and obj is not OperationHandler
+                and not inspect.isabstract(obj)
+            ):
+                # Instantiate and register the handler
+                handler = obj()
+                OPERATION_HANDLERS[handler.operation_name] = handler
+
+
+# Auto-discover and register handlers on import
+_discover_and_register_handlers()

lmcache/tools/controller_benchmark/handlers/admit.py

@@ -0,0 +1,52 @@
+"""Admit operation handler"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import TYPE_CHECKING, Any
+import random
+
+# First Party
+from lmcache.v1.cache_controller.message import (
+    BatchedKVOperationMsg,
+    KVOpEvent,
+    OpType,
+)
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+# Local
+from .base import OperationHandler
+
+
+class AdmitHandler(OperationHandler):
+    """Handler for admit operations"""
+
+    @property
+    def operation_name(self) -> str:
+        return "admit"
+
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        instance = random.choice(test_data.instances)
+        worker = random.choice(test_data.workers)
+        location = random.choice(test_data.locations)
+
+        operations = []
+        for _ in range(benchmark.config.batch_size):
+            key = random.choice(test_data.keys)
+            seq_num = benchmark.get_next_sequence_number(instance, worker, location)
+            operations.append(KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=seq_num))
+
+        return BatchedKVOperationMsg(
+            instance_id=instance,
+            worker_id=worker,
+            location=location,
+            operations=operations,
+        )
+
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        return benchmark.config.batch_size

lmcache/tools/controller_benchmark/handlers/base.py

@@ -0,0 +1,47 @@
+"""Base handler class for benchmark operations (Strategy Pattern)"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from abc import ABC, abstractmethod
+from enum import Enum, auto
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+
+class SocketType(Enum):
+    """Socket type for benchmark operations"""
+
+    PUSH = auto()  # PUSH socket for fire-and-forget messages
+    DEALER = auto()  # DEALER socket for async request-reply messages
+    HEARTBEAT = auto()  # Dedicated heartbeat DEALER socket
+
+
+class OperationHandler(ABC):
+    """Base class for benchmark operation handlers (Strategy Pattern)"""
+
+    @property
+    @abstractmethod
+    def operation_name(self) -> str:
+        """Return the operation name for registration"""
+        pass
+
+    @abstractmethod
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        """Create a message for this operation"""
+        pass
+
+    @abstractmethod
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        """Get the number of messages in a single operation"""
+        pass
+
+    @property
+    def socket_type(self) -> SocketType:
+        """Return the socket type for this operation (default: PUSH)"""
+        return SocketType.PUSH

lmcache/tools/controller_benchmark/handlers/deregister.py

@@ -0,0 +1,49 @@
+"""Deregister operation handler"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import TYPE_CHECKING, Any
+import random
+
+# First Party
+from lmcache.v1.cache_controller.message import DeRegisterMsg
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+# Local
+from .base import OperationHandler
+
+
+class DeregisterHandler(OperationHandler):
+    """Handler for deregister operations"""
+
+    @property
+    def operation_name(self) -> str:
+        return "deregister"
+
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        instance = random.choice(test_data.instances)
+        worker = random.choice(test_data.workers)
+        ip, port = self._generate_random_endpoint()
+
+        return DeRegisterMsg(
+            instance_id=instance,
+            worker_id=worker,
+            ip=ip,
+            port=port,
+        )
+
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        return 1
+
+    @staticmethod
+    def _generate_random_endpoint():
+        """Generate random IP and port for test messages"""
+        ip = "192.168.%d.%d" % (random.randint(1, 255), random.randint(1, 255))
+        port = random.randint(10000, 60000)
+        return ip, port

lmcache/tools/controller_benchmark/handlers/evict.py

@@ -0,0 +1,52 @@
+"""Evict operation handler"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import TYPE_CHECKING, Any
+import random
+
+# First Party
+from lmcache.v1.cache_controller.message import (
+    BatchedKVOperationMsg,
+    KVOpEvent,
+    OpType,
+)
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+# Local
+from .base import OperationHandler
+
+
+class EvictHandler(OperationHandler):
+    """Handler for evict operations"""
+
+    @property
+    def operation_name(self) -> str:
+        return "evict"
+
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        instance = random.choice(test_data.instances)
+        worker = random.choice(test_data.workers)
+        location = random.choice(test_data.locations)
+
+        operations = []
+        for _ in range(benchmark.config.batch_size):
+            key = random.choice(test_data.keys)
+            seq_num = benchmark.get_next_sequence_number(instance, worker, location)
+            operations.append(KVOpEvent(op_type=OpType.EVICT, key=key, seq_num=seq_num))
+
+        return BatchedKVOperationMsg(
+            instance_id=instance,
+            worker_id=worker,
+            location=location,
+            operations=operations,
+        )
+
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        return benchmark.config.batch_size

lmcache/tools/controller_benchmark/handlers/heartbeat.py

@@ -0,0 +1,56 @@
+"""Heartbeat operation handler"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import TYPE_CHECKING, Any
+import random
+
+# First Party
+from lmcache.v1.cache_controller.message import HeartbeatMsg
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+# Local
+from .base import OperationHandler, SocketType
+
+
+class HeartbeatHandler(OperationHandler):
+    """Handler for heartbeat operations (dedicated heartbeat socket)"""
+
+    @property
+    def operation_name(self) -> str:
+        return "heartbeat"
+
+    @property
+    def socket_type(self) -> SocketType:
+        return SocketType.HEARTBEAT
+
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        instance = random.choice(test_data.instances)
+        worker = random.choice(test_data.workers)
+        ip, port = self._generate_random_endpoint()
+        peer_ip, peer_port = self._generate_random_endpoint()
+        peer_init_url = "tcp://%s:%d" % (peer_ip, peer_port)
+
+        return HeartbeatMsg(
+            instance_id=instance,
+            worker_id=worker,
+            ip=ip,
+            port=port,
+            peer_init_url=peer_init_url,
+        )
+
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        return 1
+
+    @staticmethod
+    def _generate_random_endpoint():
+        """Generate random IP and port for test messages"""
+        ip = "192.168.%d.%d" % (random.randint(1, 255), random.randint(1, 255))
+        port = random.randint(10000, 60000)
+        return ip, port

lmcache/tools/controller_benchmark/handlers/p2p_lookup.py

@@ -0,0 +1,47 @@
+"""P2P lookup operation handler"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import TYPE_CHECKING, Any
+import random
+
+# First Party
+from lmcache.v1.cache_controller.message import BatchedP2PLookupMsg
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+# Local
+from .base import OperationHandler, SocketType
+
+
+class P2PLookupHandler(OperationHandler):
+    """Handler for P2P lookup operations (uses DEALER-ROUTER socket)"""
+
+    @property
+    def operation_name(self) -> str:
+        return "p2p_lookup"
+
+    @property
+    def socket_type(self) -> SocketType:
+        return SocketType.DEALER
+
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        instance = random.choice(test_data.instances)
+        worker = random.choice(test_data.workers)
+        hashes = [
+            random.choice(test_data.keys) for _ in range(benchmark.config.num_hashes)
+        ]
+
+        return BatchedP2PLookupMsg(
+            hashes=hashes,
+            instance_id=instance,
+            worker_id=worker,
+        )
+
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        return benchmark.config.num_hashes

lmcache/tools/controller_benchmark/handlers/register.py

@@ -0,0 +1,56 @@
+"""Register operation handler"""
+
+# SPDX-License-Identifier: Apache-2.0
+
+# Standard
+from typing import TYPE_CHECKING, Any
+import random
+
+# First Party
+from lmcache.v1.cache_controller.message import RegisterMsg
+
+if TYPE_CHECKING:
+    # Local
+    from ..benchmark import TestData, ZMQControllerBenchmark
+
+# Local
+from .base import OperationHandler, SocketType
+
+
+class RegisterHandler(OperationHandler):
+    """Handler for register operations (DEALER-ROUTER mode)"""
+
+    @property
+    def operation_name(self) -> str:
+        return "register"
+
+    @property
+    def socket_type(self) -> SocketType:
+        return SocketType.DEALER
+
+    def create_message(
+        self, benchmark: "ZMQControllerBenchmark", test_data: "TestData"
+    ) -> Any:
+        instance = random.choice(test_data.instances)
+        worker = random.choice(test_data.workers)
+        ip, port = self._generate_random_endpoint()
+        peer_ip, peer_port = self._generate_random_endpoint()
+        peer_init_url = "tcp://%s:%d" % (peer_ip, peer_port)
+
+        return RegisterMsg(
+            instance_id=instance,
+            worker_id=worker,
+            ip=ip,
+            port=port,
+            peer_init_url=peer_init_url,
+        )
+
+    def get_message_count(self, benchmark: "ZMQControllerBenchmark") -> int:
+        return 1
+
+    @staticmethod
+    def _generate_random_endpoint():
+        """Generate random IP and port for test messages"""
+        ip = "192.168.%d.%d" % (random.randint(1, 255), random.randint(1, 255))
+        port = random.randint(10000, 60000)
+        return ip, port

lmcache/tools/mp_status_viewer/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

lmcache/tools/mp_status_viewer/__main__.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+CLI tool to fetch and display the MP cache server status.
+
+Usage:
+    python -m lmcache.tools.mp_status_viewer [--url URL]
+"""
+
+# Standard
+import argparse
+import json
+import sys
+import urllib.error
+import urllib.request
+
+HEALTHY = "\033[32mHEALTHY\033[0m"
+UNHEALTHY = "\033[31mUNHEALTHY\033[0m"
+
+INDENT = "  "
+
+
+def _health_str(is_healthy: bool) -> str:
+    return HEALTHY if is_healthy else UNHEALTHY
+
+
+def _fmt_bytes(n: int) -> str:
+    if n >= 1024**3:
+        return f"{n / 1024**3:.2f} GB"
+    if n >= 1024**2:
+        return f"{n / 1024**2:.2f} MB"
+    if n >= 1024:
+        return f"{n / 1024:.2f} KB"
+    return f"{n} B"
+
+
+def _print_section(title: str, data: dict, depth: int = 0) -> None:
+    prefix = INDENT * depth
+    health = _health_str(data.get("is_healthy", True))
+    print(f"{prefix}{title}  [{health}]")
+
+    for key, value in data.items():
+        if key == "is_healthy":
+            continue
+        if isinstance(value, dict):
+            _print_section(key, value, depth + 1)
+        elif isinstance(value, list) and value and isinstance(value[0], dict):
+            for i, item in enumerate(value):
+                _print_section(f"{key}[{i}]", item, depth + 1)
+        else:
+            display = value
+            if "bytes" in key.lower() and isinstance(value, (int, float)):
+                display = _fmt_bytes(int(value))
+            elif isinstance(value, float):
+                display = f"{value:.4f}"
+            print(f"{prefix}{INDENT}{key}: {display}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Fetch and display LMCache MP server status"
+    )
+    parser.add_argument(
+        "--url",
+        type=str,
+        default="http://localhost:8000/api/status",
+        help="URL of the status endpoint",
+    )
+    parser.add_argument(
+        "--json",
+        action="store_true",
+        help="Output raw JSON instead of formatted view",
+    )
+    args = parser.parse_args()
+
+    try:
+        req = urllib.request.Request(args.url)
+        with urllib.request.urlopen(req, timeout=10) as resp:
+            data = json.loads(resp.read().decode())
+    except urllib.error.URLError as e:
+        print(f"Error connecting to {args.url}: {e}", file=sys.stderr)
+        sys.exit(1)
+
+    if args.json:
+        print(json.dumps(data, indent=2))
+        return
+
+    print("=" * 60)
+    print("  LMCache MP Server Status")
+    print("=" * 60)
+    _print_section("engine", data)
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()