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

lmcache/v1/mp_observability/event.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unified event model for the MP observability system."""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class EventType(Enum):
+    """All observable event types in the MP system.
+
+    Naming convention: ``<COMPONENT>_<OPERATION>`` or
+    ``<COMPONENT>_<OPERATION>_<PHASE>`` for start/end pairs.
+    """
+
+    # L1 Manager events
+    L1_READ_RESERVED = "l1.read.reserved"
+    L1_READ_FINISHED = "l1.read.finished"
+    L1_WRITE_RESERVED = "l1.write.reserved"
+    L1_WRITE_FINISHED = "l1.write.finished"
+    L1_WRITE_FINISHED_AND_READ_RESERVED = "l1.write_finished_and_read_reserved"
+    L1_KEYS_EVICTED = "l1.keys.evicted"
+
+    # StorageManager events
+    SM_READ_PREFETCHED = "sm.read.prefetched"
+    SM_READ_PREFETCHED_FINISHED = "sm.read.prefetched_finished"
+    SM_WRITE_RESERVED = "sm.write.reserved"
+    SM_WRITE_FINISHED = "sm.write.finished"
+
+    # L2 Store Controller events
+    L2_STORE_SUBMITTED = "l2.store.submitted"
+    L2_STORE_COMPLETED = "l2.store.completed"
+
+    # L2 Prefetch Controller events
+    L2_PREFETCH_LOOKUP_SUBMITTED = "l2.prefetch.lookup.submitted"
+    L2_PREFETCH_LOOKUP_COMPLETED = "l2.prefetch.lookup.completed"
+    L2_PREFETCH_LOAD_SUBMITTED = "l2.prefetch.load.submitted"
+    L2_PREFETCH_LOAD_COMPLETED = "l2.prefetch.load.completed"
+
+    # MP Server request-level events (start/end pairs)
+    MP_STORE_START = "mp.store.start"
+    MP_STORE_END = "mp.store.end"
+    MP_RETRIEVE_START = "mp.retrieve.start"
+    MP_RETRIEVE_END = "mp.retrieve.end"
+    MP_LOOKUP_PREFETCH_START = "mp.lookup_prefetch.start"
+    MP_LOOKUP_PREFETCH_END = "mp.lookup_prefetch.end"
+
+    # Chunk hash logging events
+    MP_LOOKUP = "mp.lookup"
+
+    # MP Server lifecycle sentinels (CPU-synchronous)
+    MP_REQUEST_START = "mp.request.start"
+    MP_RETRIEVE_SUBMITTED = "mp.retrieve.submitted"
+    MP_STORE_SUBMITTED = "mp.store.submitted"
+    MP_SESSION_END = "mp.session.end"
+
+    # vLLM block allocation events
+    MP_VLLM_BLOCK_ALLOCATION = "mp.vllm.block_allocation"
+
+    # vLLM end session events
+    MP_VLLM_END_SESSION = "mp.vllm.end_session"
+
+    # Trace recording — unified function-call entry event used by the
+    # ``@enable_tracing`` decorator.  Metadata layout:
+    #   ``qualname`` (str):   fully-qualified function name
+    #   ``args``     (dict):  name -> raw Python value (codec-encoded at
+    #                         record time by the recorder)
+    #   ``t_mono``   (float): ``time.monotonic()`` captured at publish
+    #                         time, so it is comparable to
+    #                         ``Event.timestamp`` (wall-clock) even
+    #                         though the drain thread processes the
+    #                         event later
+    TRACE_CALL = "trace.call"
+
+
+@dataclass
+class Event:
+    """A single observable event in the MP system.
+
+    Attributes:
+        event_type: The type of event.
+        timestamp: Wall-clock time (``time.time()``) stamped by
+            ``EventBus.publish()`` at the moment it is called — not when the
+            drain thread processes the event.  For CUDA host-callback events
+            this captures GPU-accurate timing.
+        metadata: Flat key-value payload.  Contents depend on ``event_type``;
+            see the metadata contracts in
+            ``docs/design/v1/mp_observability/event-bus.md`` Section 2.7.
+        session_id: Caller-provided ID for correlating start/end pairs.
+    """
+
+    event_type: EventType
+    timestamp: float = 0.0
+    metadata: dict[str, Any] = field(default_factory=dict)
+    session_id: str = ""

lmcache/v1/mp_observability/event_bus.py

@@ -0,0 +1,313 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""EventBus: unified pub/sub dispatcher for MP observability events."""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Callable
+import collections
+import threading
+import time
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.mp_observability.event import Event, EventType
+
+try:
+    # Third Party
+    import torch  # noqa: F401 — must be imported before lmcache.c_ops
+
+    # First Party
+    import lmcache.c_ops as _lmc_ops
+
+    _has_native_recorder = hasattr(_lmc_ops, "record_event_on_stream")
+except ImportError:
+    _has_native_recorder = False
+
+logger = init_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# Types
+# ---------------------------------------------------------------------------
+
+EventCallback = Callable[[Event], None]
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class EventBusConfig:
+    """Configuration for the EventBus.
+
+    Attributes:
+        enabled: Whether the event bus is active.  When disabled,
+            ``publish()`` is a no-op and the drain thread is not started.
+        max_queue_size: Maximum number of events in the queue.  When the
+            queue is full, new events are silently dropped with a
+            rate-limited warning.
+    """
+
+    enabled: bool = True
+    max_queue_size: int = 10_000
+
+
+# ---------------------------------------------------------------------------
+# Subscriber ABC
+# ---------------------------------------------------------------------------
+
+
+class EventSubscriber(ABC):
+    """Base class for per-component event subscribers.
+
+    Subclasses declare which ``EventType``\\s they care about via
+    ``get_subscriptions()``.  The ``register()`` helper wires them up to
+    an ``EventBus``.
+    """
+
+    @abstractmethod
+    def get_subscriptions(self) -> dict[EventType, EventCallback]:
+        """Return event_type -> callback mapping.
+
+        Called once during ``register()``.  The EventBus stores these
+        callbacks directly.
+        """
+        ...
+
+    def register(self, bus: EventBus) -> None:
+        """Subscribe all declared handlers to *bus*."""
+        for event_type, callback in self.get_subscriptions().items():
+            bus.subscribe(event_type, callback)
+
+    def shutdown(self) -> None:  # noqa: B027
+        """Optional cleanup hook.  Called by ``EventBus.stop()``."""
+        pass
+
+
+# ---------------------------------------------------------------------------
+# EventBus
+# ---------------------------------------------------------------------------
+
+
+class EventBus:
+    """Manages event ingestion, queueing, and dispatch to subscribers.
+
+    Events are appended to a deque on the hot path (``publish()``) and
+    drained by a background thread that dispatches to registered callbacks.
+    """
+
+    def __init__(self, config: EventBusConfig | None = None) -> None:
+        if config is None:
+            config = EventBusConfig()
+        self._config = config
+        self._subscribers: dict[EventType, list[EventCallback]] = defaultdict(list)
+        self._queue: collections.deque[Event] = collections.deque()
+        self._wake = threading.Event()
+        self._stop_flag = threading.Event()
+        self._lock = threading.Lock()
+        self._thread: threading.Thread | None = None
+        self._registered_subscribers: list[EventSubscriber] = []
+        self._discard_count: int = 0
+        self._last_discard_warning: float = 0.0
+
+    # -- Public API --------------------------------------------------------
+
+    def subscribe(self, event_type: EventType, callback: EventCallback) -> None:
+        """Register a callback for a specific event type (thread-safe)."""
+        with self._lock:
+            self._subscribers[event_type].append(callback)
+
+    def register_subscriber(self, subscriber: EventSubscriber) -> None:
+        """Register an ``EventSubscriber`` and wire up its callbacks."""
+        subscriber.register(self)
+        with self._lock:
+            self._registered_subscribers.append(subscriber)
+
+    def has_subscribers(self, event_type: EventType) -> bool:
+        """Return True if at least one callback is registered for *event_type*.
+
+        Use this to skip expensive event construction on the hot path when
+        no subscriber is listening::
+
+            if bus.has_subscribers(EventType.MP_LOOKUP):
+                bus.publish(Event(event_type=EventType.MP_LOOKUP, ...))
+        """
+        return bool(self._subscribers.get(event_type))
+
+    def publish_on_stream(self, stream: Any, event: Event) -> None:
+        """Schedule event recording as a CUDA host function on *stream*.
+
+        Uses a C++ callback via ``cudaLaunchHostFunc`` so the callback
+        never touches the GIL, avoiding the CUDA-driver/GIL deadlock.
+
+        No-op when the EventBus is disabled, avoiding the overhead of
+        scheduling a host function on the CUDA stream entirely.
+        """
+        if not self._config.enabled:
+            return
+        if _has_native_recorder:
+            str_metadata: dict[str, str] = {}
+            int_metadata: dict[str, int] = {}
+            for k, v in event.metadata.items():
+                if isinstance(v, int):
+                    int_metadata[k] = v
+                else:
+                    str_metadata[k] = str(v)
+            _lmc_ops.record_event_on_stream(
+                stream.ptr,
+                event.event_type.value,
+                event.session_id,
+                str_metadata,
+                int_metadata,
+            )
+        else:
+            stream.launch_host_func(self.publish, event)
+
+    def publish(self, event: Event) -> None:
+        """Submit an event (hot path — non-blocking).
+
+        The event's ``timestamp`` is set to ``time.time()`` at call time.
+        When the queue is full the event is silently discarded with a
+        rate-limited warning (at most once per second).
+        """
+        if not self._config.enabled:
+            return
+
+        if len(self._queue) >= self._config.max_queue_size:
+            self._discard_count += 1
+            now = time.monotonic()
+            if now - self._last_discard_warning >= 1.0:
+                logger.warning(
+                    "EventBus queue full (max_queue_size=%d), "
+                    "%d event(s) discarded so far",
+                    self._config.max_queue_size,
+                    self._discard_count,
+                )
+                self._last_discard_warning = now
+            return
+
+        event.timestamp = time.time()
+        self._queue.append(event)
+        self._wake.set()
+
+    def start(self) -> None:
+        """Start the background drain thread.  No-op when disabled or
+        already running."""
+        if not self._config.enabled:
+            return
+        if self._thread is not None and self._thread.is_alive():
+            return
+
+        self._stop_flag.clear()
+        self._thread = threading.Thread(
+            target=self._run,
+            daemon=True,
+            name="EventBus",
+        )
+        logger.debug("Starting EventBus drain thread...")
+        self._thread.start()
+
+    def stop(self) -> None:
+        """Stop the drain thread, flush remaining events, and shut down
+        all registered subscribers.  Safe to call when not started."""
+        self._stop_flag.set()
+        self._wake.set()
+        if self._thread is not None and self._thread.is_alive():
+            self._thread.join()
+
+        # Final drain
+        self._drain_all()
+
+        # Shutdown subscribers
+        with self._lock:
+            snapshot = list(self._registered_subscribers)
+        for sub in snapshot:
+            try:
+                sub.shutdown()
+            except Exception:
+                logger.exception(
+                    "EventBus: error shutting down %s",
+                    type(sub).__name__,
+                )
+
+    # -- Internal ----------------------------------------------------------
+
+    def _run(self) -> None:
+        """Drain loop: wait for wake signal or timeout, then drain."""
+        while not self._stop_flag.is_set():
+            self._wake.wait(timeout=0.1)
+            self._wake.clear()
+            self._drain_all()
+
+    def _drain_all(self) -> None:
+        """Pop all queued events and dispatch to subscribers."""
+        # Drain events buffered on the C++ side (from CUDA host callbacks)
+        if _has_native_recorder:
+            for name, sid, ts, str_meta, int_meta in _lmc_ops.drain_recorded_events():
+                metadata: dict[str, Any] = dict(str_meta)
+                metadata.update(int_meta)
+                self._queue.append(
+                    Event(
+                        event_type=EventType(name),
+                        session_id=sid,
+                        timestamp=ts,
+                        metadata=metadata,
+                    )
+                )
+
+        with self._lock:
+            snapshot = dict(self._subscribers)
+
+        while True:
+            try:
+                event = self._queue.popleft()
+            except IndexError:
+                break
+            for cb in snapshot.get(event.event_type, []):
+                try:
+                    cb(event)
+                except Exception:
+                    logger.exception(
+                        "EventBus: error in callback for %s",
+                        event.event_type.value,
+                    )
+
+
+# ---------------------------------------------------------------------------
+# Singleton
+# ---------------------------------------------------------------------------
+
+_global_bus = EventBus(EventBusConfig(enabled=False))
+_observability_enabled: bool = False
+
+
+def is_observability_enabled() -> bool:
+    """Fast check for whether observability is active.
+
+    Use this to guard expensive event-construction or CUDA host-function
+    scheduling when observability is disabled.
+    """
+    return _observability_enabled
+
+
+def get_event_bus() -> EventBus:
+    """Return the current global EventBus singleton."""
+    return _global_bus
+
+
+def init_event_bus(config: EventBusConfig | None = None) -> EventBus:
+    """Replace the global singleton with a new EventBus built from *config*.
+
+    Returns the newly created bus.
+    """
+    global _global_bus, _observability_enabled
+    _global_bus = EventBus(config)
+    _observability_enabled = config.enabled if config else True
+    return _global_bus

lmcache/v1/mp_observability/otel_init.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""OpenTelemetry SDK initialization for the MP observability system.
+
+Supports two modes, controlled by the ``otlp_endpoint`` field in
+``ObservabilityConfig``:
+
+- **OTLP push** (production): metrics/traces are pushed to an OTel collector.
+- **Prometheus pull** (dev/debug): metrics are served on a local ``/metrics``
+  endpoint via ``prometheus_client``, no collector needed.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from collections.abc import Callable
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+def init_otel_metrics(
+    otlp_endpoint: str | None = None,
+    prometheus_port: int | None = None,
+) -> None:
+    """Set up the OpenTelemetry MeterProvider.
+
+    Args:
+        otlp_endpoint: OTLP gRPC endpoint (e.g. ``http://localhost:4317``).
+            When set, metrics are pushed to an OTel collector.
+            When ``None``, falls back to Prometheus pull mode.
+        prometheus_port: Port for the fallback Prometheus ``/metrics``
+            endpoint.  Only used when *otlp_endpoint* is ``None``.
+            Defaults to 9090.
+    """
+    # Third Party
+    from opentelemetry import metrics
+    from opentelemetry.sdk.metrics import MeterProvider
+
+    if otlp_endpoint is not None:
+        # OTLP push mode
+        # Third Party
+        from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import (
+            OTLPMetricExporter,
+        )
+        from opentelemetry.sdk.metrics.export import (
+            PeriodicExportingMetricReader,
+        )
+
+        exporter = OTLPMetricExporter(endpoint=otlp_endpoint, insecure=True)
+        reader = PeriodicExportingMetricReader(exporter, export_interval_millis=10000)
+        provider = MeterProvider(metric_readers=[reader])
+        metrics.set_meter_provider(provider)
+        logger.info(
+            "OTel MeterProvider initialised with OTLP exporter (%s)",
+            otlp_endpoint,
+        )
+    else:
+        # Prometheus pull fallback — no collector needed
+        # Third Party
+        from opentelemetry.exporter.prometheus import PrometheusMetricReader
+        import prometheus_client
+
+        if prometheus_port is None:
+            prometheus_port = 9090
+
+        reader = PrometheusMetricReader()
+        provider = MeterProvider(metric_readers=[reader])
+        metrics.set_meter_provider(provider)
+        prometheus_client.start_http_server(prometheus_port)
+        logger.info(
+            "OTel MeterProvider initialised with Prometheus fallback "
+            "(http://0.0.0.0:%d/metrics)",
+            prometheus_port,
+        )
+
+
+def init_otel_tracing(otlp_endpoint: str | None = None) -> None:
+    """Set up the OpenTelemetry TracerProvider with an OTLP exporter.
+
+    Tracing requires an OTLP endpoint — there is no local fallback.
+    When *otlp_endpoint* is ``None``, tracing init is skipped.
+
+    Args:
+        otlp_endpoint: OTLP gRPC endpoint.  When ``None``, tracing
+            init is skipped (no-op).
+    """
+    if otlp_endpoint is None:
+        logger.debug("No OTLP endpoint configured, skipping tracing init")
+        return
+
+    # Third Party
+    from opentelemetry import trace
+    from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+        OTLPSpanExporter,
+    )
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+    exporter = OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True)
+    provider = TracerProvider()
+    provider.add_span_processor(BatchSpanProcessor(exporter))
+    trace.set_tracer_provider(provider)
+    logger.info(
+        "OTel TracerProvider initialised with OTLP exporter (%s)",
+        otlp_endpoint,
+    )
+
+
+def register_gauge(
+    meter_name: str,
+    gauge_name: str,
+    description: str,
+    func: Callable[[], int | float],
+) -> None:
+    """Register an OTel observable gauge with a callback.
+
+    This is a convenience wrapper that hides the OTel boilerplate.
+    If OTel is not available, the call is silently ignored.
+
+    Args:
+        meter_name: OTel meter name (e.g. ``lmcache.mp_engine``).
+        gauge_name: Metric name (e.g.
+            ``lmcache_mp.active_prefetch_jobs``).
+        description: Human-readable description of the gauge.
+        func: Zero-arg callable returning the current value.
+    """
+    try:
+        # Third Party
+        from opentelemetry import metrics as otel_metrics
+
+        meter = otel_metrics.get_meter(meter_name)
+        meter.create_observable_gauge(
+            gauge_name,
+            callbacks=[lambda _: [otel_metrics.Observation(func())]],
+            description=description,
+        )
+    except ImportError:
+        logger.debug(
+            "opentelemetry package not found, skipping gauge %s",
+            gauge_name,
+        )

lmcache/v1/mp_observability/subscribers/__init__.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# First Party
+from lmcache.v1.mp_observability.subscribers.logging import (
+    L1LoggingSubscriber,
+    MPServerLoggingSubscriber,
+    SMLoggingSubscriber,
+)
+from lmcache.v1.mp_observability.subscribers.metrics import (
+    L0LifecycleSubscriber,
+    L1LifecycleSubscriber,
+    L1MetricsSubscriber,
+    SMMetricsSubscriber,
+)
+from lmcache.v1.mp_observability.subscribers.tracing import (
+    MPServerTracingSubscriber,
+)
+
+__all__ = [
+    "L0LifecycleSubscriber",
+    "L1LifecycleSubscriber",
+    "L1LoggingSubscriber",
+    "L1MetricsSubscriber",
+    "MPServerLoggingSubscriber",
+    "MPServerTracingSubscriber",
+    "SMLoggingSubscriber",
+    "SMMetricsSubscriber",
+]

lmcache/v1/mp_observability/subscribers/logging/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# First Party
+from lmcache.v1.mp_observability.subscribers.logging.l1 import L1LoggingSubscriber
+from lmcache.v1.mp_observability.subscribers.logging.l2 import L2LoggingSubscriber
+from lmcache.v1.mp_observability.subscribers.logging.lookup_hash import (
+    LookupHashLoggingSubscriber,
+)
+from lmcache.v1.mp_observability.subscribers.logging.mp_server import (
+    MPServerLoggingSubscriber,
+)
+from lmcache.v1.mp_observability.subscribers.logging.sm import SMLoggingSubscriber
+
+__all__ = [
+    "LookupHashLoggingSubscriber",
+    "L1LoggingSubscriber",
+    "L2LoggingSubscriber",
+    "MPServerLoggingSubscriber",
+    "SMLoggingSubscriber",
+]

lmcache/v1/mp_observability/subscribers/logging/l1.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""L1 logging subscriber — debug logs for L1Manager events.
+
+Logs are emitted via Python's standard logging module.  When OpenTelemetry
+is installed, ``init_logger`` automatically attaches an OTel
+``LoggingHandler`` so records are forwarded to OTel when a
+``LoggerProvider`` is configured at startup.
+"""
+
+# Future
+from __future__ import annotations
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.mp_observability.event import Event, EventType
+from lmcache.v1.mp_observability.event_bus import EventCallback, EventSubscriber
+
+logger = init_logger(__name__)
+
+
+class L1LoggingSubscriber(EventSubscriber):
+    """Logs L1Manager events at debug level."""
+
+    def get_subscriptions(self) -> dict[EventType, EventCallback]:
+        return {
+            EventType.L1_READ_RESERVED: self._on_read_reserved,
+            EventType.L1_READ_FINISHED: self._on_read_finished,
+            EventType.L1_WRITE_RESERVED: self._on_write_reserved,
+            EventType.L1_WRITE_FINISHED: self._on_write_finished,
+            EventType.L1_WRITE_FINISHED_AND_READ_RESERVED: (
+                self._on_write_finished_and_read_reserved
+            ),
+            EventType.L1_KEYS_EVICTED: self._on_evicted,
+        }
+
+    def _on_read_reserved(self, event: Event) -> None:
+        logger.debug("L1 read reserved: %d keys", len(event.metadata["keys"]))
+
+    def _on_read_finished(self, event: Event) -> None:
+        logger.debug("L1 read finished: %d keys", len(event.metadata["keys"]))
+
+    def _on_write_reserved(self, event: Event) -> None:
+        logger.debug("L1 write reserved: %d keys", len(event.metadata["keys"]))
+
+    def _on_write_finished(self, event: Event) -> None:
+        logger.debug("L1 write finished: %d keys", len(event.metadata["keys"]))
+
+    def _on_write_finished_and_read_reserved(self, event: Event) -> None:
+        logger.debug(
+            "L1 write finished and read reserved: %d keys",
+            len(event.metadata["keys"]),
+        )
+
+    def _on_evicted(self, event: Event) -> None:
+        logger.debug("L1 eviction: %d keys", len(event.metadata["keys"]))