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

lmcache/v1/cache_controller/observability.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from enum import Enum
+from typing import Optional
+
+# Third Party
+from prometheus_client import REGISTRY
+import prometheus_client
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class SocketType(Enum):
+    """Enum for socket types to ensure type safety."""
+
+    PULL = "pull"
+    REPLY = "reply"
+
+
+class PrometheusLogger:
+    """
+    Prometheus logger for cache controller metrics.
+    Provides dynamic metrics for monitoring KV pool and worker registration.
+    """
+
+    _instance = None
+    _gauge_cls = prometheus_client.Gauge
+
+    def __init__(self, labels: dict):
+        self.labels = labels
+        labelnames = list(self.labels.keys())
+
+        # Dynamic metrics for cache controller
+        self._init_dynamic_metrics(labelnames)
+
+    def _init_dynamic_metrics(self, labelnames):
+        """
+        Initialize dynamic metrics that will be updated by lambda functions.
+        """
+        # KV Pool metrics
+        self.kv_pool_keys_count = self._gauge_cls(
+            name="lmcache:cache_controller_kv_pool_keys_count",
+            documentation="The number of keys in the KV pool",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        # Registration Controller metrics
+        self.registered_workers_count = self._gauge_cls(
+            name="lmcache:cache_controller_registered_workers_count",
+            documentation="The total number of registered workers",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        # Socket message count metrics
+        self.pull_socket_message_count = self._gauge_cls(
+            name="lmcache:cache_controller_pull_socket_message_count",
+            documentation="The total number of messages received on PULL socket",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        self.reply_socket_message_count = self._gauge_cls(
+            name="lmcache:cache_controller_reply_socket_message_count",
+            documentation="The total number of messages received on REPLY socket",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+
+        # Socket queue/backlog metrics
+        self.pull_socket_has_pending = self._gauge_cls(
+            name="lmcache:cache_controller_pull_socket_has_pending",
+            documentation="Whether PULL socket has pending messages (1=yes, 0=no)",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        self.reply_socket_has_pending = self._gauge_cls(
+            name="lmcache:cache_controller_reply_socket_has_pending",
+            documentation="Whether REPLY socket has pending messages (1=yes, 0=no)",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+
+        # Active request metrics
+        self.pull_socket_active_requests = self._gauge_cls(
+            name="lmcache:cache_controller_pull_socket_active_requests",
+            documentation="Number of requests being processed from PULL socket",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        self.reply_socket_active_requests = self._gauge_cls(
+            name="lmcache:cache_controller_reply_socket_active_requests",
+            documentation="Number of requests being processed from REPLY socket",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+
+        # Sequence number discontinuity metrics
+        self.kv_op_seq_discontinuity_count = self._gauge_cls(
+            name="lmcache:cache_controller_kv_op_seq_discontinuity_count",
+            documentation="Total count of KV operation sequence number discontinuities",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+
+        # Full sync metrics
+        self.full_sync_workers_syncing = self._gauge_cls(
+            name="lmcache:cache_controller_full_sync_workers_syncing",
+            documentation="Number of workers currently in full sync",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        self.full_sync_workers_completed = self._gauge_cls(
+            name="lmcache:cache_controller_full_sync_workers_completed",
+            documentation="Number of workers that have completed full sync",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        self.full_sync_global_progress = self._gauge_cls(
+            name="lmcache:cache_controller_full_sync_global_progress",
+            documentation="Global full sync progress (0.0 to 1.0)",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+        self.full_sync_missing_batches_total = self._gauge_cls(
+            name="lmcache:cache_controller_full_sync_missing_batches_total",
+            documentation="Total count of missing batches across all syncing workers",
+            labelnames=labelnames,
+            multiprocess_mode="livemostrecent",
+        ).labels(**self.labels)
+
+    @staticmethod
+    def GetOrCreate(
+        labels: dict,
+    ) -> "PrometheusLogger":
+        if PrometheusLogger._instance is None:
+            PrometheusLogger._instance = PrometheusLogger(labels)
+        if PrometheusLogger._instance.labels != labels:
+            logger.error(
+                "CacheControllerPrometheusLogger instance already created with "
+                "different metadata. This should not happen except in test"
+            )
+        return PrometheusLogger._instance
+
+    @staticmethod
+    def GetInstance() -> "PrometheusLogger":
+        assert PrometheusLogger._instance is not None, (
+            "CacheControllerPrometheusLogger instance not created yet"
+        )
+        return PrometheusLogger._instance
+
+    @staticmethod
+    def GetInstanceOrNone() -> Optional["PrometheusLogger"]:
+        """
+        Returns the singleton instance of CacheControllerPrometheusLogger if it exists,
+        otherwise returns None.
+        """
+        return PrometheusLogger._instance
+
+    @staticmethod
+    def DestroyInstance():
+        PrometheusLogger._instance = None
+
+    @staticmethod
+    def unregister_all_metrics():
+        collectors = list(REGISTRY._collector_to_names.keys())
+        for collector in collectors:
+            try:
+                REGISTRY.unregister(collector)
+            except KeyError:
+                pass
+
+
+class SocketMetricsContext:
+    """Context manager for socket message counting and error handling."""
+
+    def __init__(self, manager, socket_type: SocketType, message_count: int = 1):
+        self.manager = manager
+        self.socket_type = socket_type
+        self.message_count = message_count
+        self.counter_attr = f"{socket_type.value}_socket_message_count"
+        self.active_attr = f"{socket_type.value}_socket_active_requests"
+
+    def __enter__(self):
+        setattr(
+            self.manager,
+            self.counter_attr,
+            getattr(self.manager, self.counter_attr) + self.message_count,
+        )
+        setattr(
+            self.manager,
+            self.active_attr,
+            getattr(self.manager, self.active_attr) + self.message_count,
+        )
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        setattr(
+            self.manager,
+            self.active_attr,
+            getattr(self.manager, self.active_attr) - self.message_count,
+        )
+        if exc_type is not None:
+            logger.error(
+                "Controller Manager error", exc_info=(exc_type, exc_val, exc_tb)
+            )
+        return False