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

lmcache/v1/metadata.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass, field
+from typing import Optional
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.kv_layer_groups import KVLayerGroupsManager
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class LMCacheMetadata:
+    """
+    LMCacheMetadata should be extracted from the northbound
+    serving engine configuration and wrap the extraction of
+    attributes (e.g. model name, tp rank, etc.)
+    """
+
+    """name of the LLM model"""
+    model_name: str
+    """ global world size when running under a distributed setting 
+    (total number of workers)"""
+    world_size: int
+    """ host world size (workers on active localhost)
+    This information can be useful for multi-node
+    deployment. Will be the same as world_size 
+    in single-node deployments.
+    """
+    local_world_size: int
+    """ worker id when running under a distributed setting """
+    worker_id: int
+    """ host worker id (a gpu bound worker id on active localhost)
+    This information can be useful for multi-node deployment. 
+    Will be the same as worker_id in single-node deployments.
+    """
+    local_worker_id: int
+    """ the data type of kv tensors """
+    # (Deprecated) Will be replaced by kv_layer_groups_manager in the future
+    kv_dtype: torch.dtype
+    """ the shape of kv tensors """
+    # (Deprecated) Will be replaced by kv_layer_groups_manager in the future
+    """ (num_layer, 2, chunk_size, num_kv_head, head_size) """
+    kv_shape: tuple[int, int, int, int, int]
+    """ whether use MLA"""
+    use_mla: bool = False
+    """ the role of the current instance (e.g., 'scheduler', 'worker') """
+    role: Optional[str] = None
+    """ the first rank of the distributed setting """
+    # TODO(baoloongmao): first_rank should be configurable
+    first_rank = 0
+    served_model_name: Optional[str] = None
+    """chunk size"""
+    chunk_size: int = 256
+    """ Manager for groups of layers with identical KV cache structure """
+    kv_layer_groups_manager: KVLayerGroupsManager = field(
+        default_factory=KVLayerGroupsManager
+    )
+    """ engine_id for RPC path (used by lookup client/server) """
+    engine_id: Optional[str] = None
+    """ extra config from kv_connector (e.g., lmcache_rpc_port) """
+    kv_connector_extra_config: Optional[dict] = None
+
+    def is_first_rank(self) -> bool:
+        """Check if the current worker is the first rank"""
+        return self.worker_id == self.first_rank
+
+    # TODO(chunxiaozheng): some uts do not `build_kv_layer_groups`
+    def get_dtypes(self) -> list[torch.dtype]:
+        if self.kv_layer_groups_manager.kv_layer_groups:
+            return [
+                group.dtype for group in self.kv_layer_groups_manager.kv_layer_groups
+            ]
+        return [self.kv_dtype]
+
+    def get_shapes(self, num_tokens: Optional[int] = None) -> list[torch.Size]:
+        """Get the shapes of the KV cache in LMCache"""
+        if num_tokens is None:
+            num_tokens = self.chunk_size
+        if self.kv_layer_groups_manager.kv_layer_groups:
+            shapes = []
+            kv_size = 1 if self.use_mla else 2
+            for group in self.kv_layer_groups_manager.kv_layer_groups:
+                shapes.append(
+                    torch.Size(
+                        [
+                            kv_size,
+                            group.num_layers,
+                            num_tokens,
+                            group.hidden_dim_size,
+                        ]
+                    )
+                )
+            return shapes
+        else:
+            return [
+                torch.Size(
+                    [
+                        self.kv_shape[1],
+                        self.kv_shape[0],
+                        num_tokens,
+                        self.kv_shape[3] * self.kv_shape[4],
+                    ]
+                )
+            ]
+
+    def get_num_groups(self) -> int:
+        if self.kv_layer_groups_manager.kv_layer_groups:
+            return self.kv_layer_groups_manager.num_groups
+        return 1

lmcache/v1/mp_observability/AGENTS.override.md

@@ -0,0 +1,21 @@
+# MP Observability — Agent Instructions
+
+When working in `lmcache/v1/mp_observability/`:
+
+Design docs live in `docs/design/v1/mp_observability/` — update them whenever
+the contracts below change.
+
+1. **New `EventType`** — after adding an entry to `event.py`, update the
+   metadata contract table in `docs/design/v1/mp_observability/EVENTS.md`
+   with the new type's metadata keys and types.
+
+2. **New metrics subscriber** — after adding counters/histograms, update the
+   metrics table in `docs/design/v1/mp_observability/METRICS.md` with the
+   metric name, type, and description.
+
+3. **New subscriber class** — follow the step-by-step guide in `README.md`
+   (co-located with this file; "How to Add a New Event and Subscriber") and
+   the design rules in `docs/design/v1/mp_observability/event-bus.md`.
+
+4. **CLI args** — if you add or change observability CLI flags in `config.py`,
+   update `docs/source/mp/observability.rst` to match.

lmcache/v1/mp_observability/README.md

@@ -0,0 +1,204 @@
+# MP Observability
+
+Event-driven observability for LMCache's multiprocess (MP) mode, built on
+[OpenTelemetry](https://opentelemetry.io/).
+
+For metrics, see [METRICS.md](../../../docs/design/v1/mp_observability/METRICS.md).
+For event metadata contracts, see
+[EVENTS.md](../../../docs/design/v1/mp_observability/EVENTS.md).
+For design rationale, see
+[event-bus.md](../../../docs/design/v1/mp_observability/event-bus.md).
+For the trace recording subsystem (`lmcache trace`), see
+[trace.md](../../../docs/design/v1/mp_observability/trace.md).
+
+---
+
+## Architecture
+
+```
+Producers (L1Manager, StorageManager, MPCacheEngine)
+    │
+    │  event_bus.publish(Event(...))
+    ▼
+EventBus  (async queue + drain thread)
+    │
+    ├──► L1MetricsSubscriber          → OTel counter.add(...)
+    ├──► SMMetricsSubscriber          → OTel counter.add(...)
+    ├──► L1LoggingSubscriber          → logger.debug(...)
+    ├──► SMLoggingSubscriber          → logger.debug(...)
+    ├──► MPServerLoggingSubscriber    → logger.debug(...)
+    └──► MPServerTracingSubscriber    → OTel span start/end
+
+OTel SDK  (configured at startup)
+    │
+    ├──► OTLP push (production)       → OTel collector → Prometheus / Grafana / etc.
+    └──► Prometheus pull (dev/debug)   → /metrics on configured port
+```
+
+---
+
+## Configuration
+
+All observability behaviour is controlled by `ObservabilityConfig`
+(defined in `config.py`).  When running the LMCache MP mode server from the
+CLI, pass the flags below; when embedding programmatically, construct an
+`ObservabilityConfig` directly.
+
+### CLI flags
+
+| Flag | Default | Description |
+|---|---|---|
+| `--disable-observability` | off | Disable the EventBus entirely. No events are published or consumed. |
+| `--disable-metrics` | off | Skip registering metrics subscribers (OTel counters). |
+| `--disable-logging` | off | Skip registering logging subscribers. |
+| `--enable-tracing` | off | Register tracing subscribers (OTel spans). Disabled by default. **Requires `--otlp-endpoint`.** |
+| `--event-bus-queue-size N` | `10000` | Maximum number of events in the EventBus queue before tail-drop. |
+| `--otlp-endpoint URL` | *(none)* | OTLP gRPC endpoint (e.g. `http://localhost:4317`). When set, metrics and traces are pushed to an OTel collector. When unset, metrics fall back to Prometheus pull mode. |
+| `--prometheus-port PORT` | `9090` | Port for the Prometheus `/metrics` endpoint. Only used when `--otlp-endpoint` is not set. |
+
+### `ObservabilityConfig` fields
+
+| Field | Type | Default | Description |
+|---|---|---|---|
+| `enabled` | `bool` | `True` | Master switch for the EventBus. |
+| `max_queue_size` | `int` | `10000` | Maximum events in the EventBus queue before tail-drop. |
+| `metrics_enabled` | `bool` | `True` | Register metrics subscribers (OTel counters / histograms). |
+| `logging_enabled` | `bool` | `True` | Register logging subscribers. |
+| `tracing_enabled` | `bool` | `False` | Register tracing subscribers (OTel spans). |
+| `otlp_endpoint` | `str \| None` | `None` | OTLP gRPC endpoint. When set, metrics and traces are pushed. When `None`, metrics use Prometheus pull fallback. |
+| `prometheus_port` | `int` | `9090` | Port for the Prometheus `/metrics` endpoint (pull fallback only). |
+
+### Metrics export modes
+
+| `otlp_endpoint` | Mode | How to query |
+|---|---|---|
+| `http://host:4317` | OTLP push | Query the OTel collector's Prometheus exporter |
+| `None` | Prometheus pull fallback | `curl http://localhost:<prometheus-port>/metrics` |
+
+> **Note:** OTel counters only appear on `/metrics` after the first increment.
+> If you see only Python runtime metrics, trigger a store/retrieve first.
+
+### Tracing
+
+Tracing is opt-in (`--enable-tracing`).  When enabled, `MPServerTracingSubscriber`
+creates OTel spans from MP server START/END event pairs (store, retrieve,
+lookup/prefetch).  Trace export requires an OTLP endpoint — there is no local
+fallback.  `--enable-tracing` **requires** `--otlp-endpoint`; the server will
+raise a `ValueError` at startup if the endpoint is missing.
+
+---
+
+## How to Add a New Event and Subscriber
+
+### Step 1 — Define the event type
+
+Add a new member to `EventType` in `event.py`:
+
+```python
+class EventType(Enum):
+    # ... existing events ...
+
+    # My new component events
+    MY_COMPONENT_OPERATION = "my_component.operation"
+```
+
+### Step 2 — Publish the event from the producer
+
+In your component (e.g., a manager class), publish to the EventBus:
+
+```python
+from lmcache.v1.mp_observability.event import Event, EventType
+from lmcache.v1.mp_observability.event_bus import get_event_bus
+
+class MyComponent:
+    def __init__(self):
+        self._event_bus = get_event_bus()
+
+    def do_operation(self, keys):
+        # ... business logic ...
+
+        self._event_bus.publish(Event(
+            event_type=EventType.MY_COMPONENT_OPERATION,
+            metadata={"keys": keys},
+        ))
+```
+
+### Step 3 — Create a subscriber
+
+Create a file under the appropriate `subscribers/` subdirectory:
+
+- `subscribers/metrics/` for OTel counters / histograms
+- `subscribers/logging/` for debug log output
+- `subscribers/tracing/` for OTel spans
+
+Example metrics subscriber (`subscribers/metrics/my_component.py`):
+
+```python
+from opentelemetry import metrics
+from lmcache.v1.mp_observability.event import Event, EventType
+from lmcache.v1.mp_observability.event_bus import EventCallback, EventSubscriber
+
+
+class MyComponentMetricsSubscriber(EventSubscriber):
+    def __init__(self):
+        meter = metrics.get_meter("lmcache.my_component")
+        self._op_counter = meter.create_counter(
+            "lmcache_mp.my_component_operations",
+            description="Total operations on my component",
+        )
+
+    def get_subscriptions(self) -> dict[EventType, EventCallback]:
+        return {
+            EventType.MY_COMPONENT_OPERATION: self._on_operation,
+        }
+
+    def _on_operation(self, event: Event) -> None:
+        self._op_counter.add(len(event.metadata["keys"]))
+```
+
+### Step 4 — Export from `__init__.py`
+
+Add the subscriber to the corresponding `__init__.py` so it can be
+imported from the package:
+
+```python
+# subscribers/metrics/__init__.py
+from lmcache.v1.mp_observability.subscribers.metrics.my_component import (
+    MyComponentMetricsSubscriber,
+)
+```
+
+### Step 5 — Register the subscriber at startup
+
+In the server startup function (e.g., `run_cache_server()` in `server.py`),
+register conditionally based on `ObservabilityConfig`:
+
+```python
+if obs_config.metrics_enabled:
+    from lmcache.v1.mp_observability.subscribers.metrics import (
+        MyComponentMetricsSubscriber,
+    )
+    bus.register_subscriber(MyComponentMetricsSubscriber())
+```
+
+### Step 6 — Document the metadata contract
+
+Add a row to the metadata contracts table in
+[EVENTS.md](../../../docs/design/v1/mp_observability/EVENTS.md) so
+subscribers can rely on the schema:
+
+```markdown
+| `MY_COMPONENT_OPERATION` | `keys` | `list[ObjectKey]` |
+```
+
+---
+
+## Design rules
+
+| Rule | Reason |
+|---|---|
+| Create meters and counters in `__init__()`, not at module level | `MeterProvider` must be set before `get_meter()` is called. Module-level calls happen at import time, before setup. |
+| Prefix OTel metric names with `lmcache_mp.` | Keeps the MP namespace separate from `lmcache.` (the single-process engine namespace). |
+| Use `metadata: dict[str, Any]` for event payloads | Flexible, no coupling between producers and subscribers. See metadata contracts in [EVENTS.md](../../../docs/design/v1/mp_observability/EVENTS.md). |
+| Separate metrics, logging, and tracing subscribers | Single responsibility. Can enable/disable independently via config. |
+| Store `self._event_bus = get_event_bus()` in `__init__` | Avoids calling the singleton getter on every publish. |

lmcache/v1/mp_observability/config.py

@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Configuration for the MP-mode observability stack.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+import argparse
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.mp_observability.event_bus import EventBus
+
+# First Party
+from lmcache.v1.mp_observability.subscribers.logging.lookup_hash import (
+    LookupHashLogConfig,
+)
+
+
+@dataclass
+class ObservabilityConfig:
+    """Unified configuration for the EventBus-based observability system.
+
+    Controls the EventBus, OTel metrics/tracing pipelines, and subscriber
+    registration.
+    """
+
+    enabled: bool = True
+    """Master switch for the EventBus."""
+
+    max_queue_size: int = 10_000
+    """Maximum events in the EventBus queue before tail-drop."""
+
+    metrics_enabled: bool = True
+    """Register metrics subscribers (OTel counters / histograms)."""
+
+    logging_enabled: bool = True
+    """Register logging subscribers."""
+
+    tracing_enabled: bool = False
+    """Register span subscribers (OTel traces)."""
+
+    otlp_endpoint: str | None = None
+    """OTLP gRPC endpoint (e.g. ``http://localhost:4317``).  When set,
+    metrics and traces are pushed to an OTel collector.  When ``None``,
+    metrics fall back to an in-process Prometheus ``/metrics`` endpoint."""
+
+    prometheus_port: int = 9090
+    """Port for the Prometheus /metrics endpoint.  Only used when
+    ``otlp_endpoint`` is ``None`` (Prometheus pull fallback)."""
+
+    metrics_sample_rate: float = 0.01
+    """Fraction of chunks/blocks to track for lifecycle histograms (0, 1.0].
+    Counters always count all events regardless of this setting."""
+
+    lookup_hash_log: LookupHashLogConfig = field(default_factory=LookupHashLogConfig)
+    """Configuration for lookup hash file logging.  Disabled by default
+    (empty ``output_dir``)."""
+
+    trace_level: str | None = None
+    """If set, enables trace recording at the given level.  Currently
+    only ``"storage"`` is supported.  See
+    :mod:`lmcache.v1.mp_observability.trace` for details."""
+
+    trace_output: str | None = None
+    """Path to write the trace file.  When :attr:`trace_level` is set
+    but this is ``None``, a timestamped path under ``$TMPDIR`` is
+    minted and logged at INFO."""
+
+
+DEFAULT_OBSERVABILITY_CONFIG = ObservabilityConfig(enabled=False)
+
+
+def add_observability_args(
+    parser: argparse.ArgumentParser,
+) -> argparse.ArgumentParser:
+    """Add observability configuration arguments to an existing parser.
+
+    Args:
+        parser: The argument parser to add arguments to.
+
+    Returns:
+        The same parser with observability arguments added.
+    """
+    group = parser.add_argument_group(
+        "Observability", "Configuration for metrics, logging, and tracing"
+    )
+    group.add_argument(
+        "--disable-observability",
+        action="store_true",
+        default=False,
+        help="Disable the observability EventBus entirely.",
+    )
+    group.add_argument(
+        "--disable-metrics",
+        action="store_true",
+        default=False,
+        help="Disable metrics subscribers (OTel counters).",
+    )
+    group.add_argument(
+        "--disable-logging",
+        action="store_true",
+        default=False,
+        help="Disable logging subscribers.",
+    )
+    group.add_argument(
+        "--enable-tracing",
+        action="store_true",
+        default=False,
+        help="Enable span subscribers (OTel traces). Disabled by default.",
+    )
+    group.add_argument(
+        "--otlp-endpoint",
+        type=str,
+        default=None,
+        help=(
+            "OTLP gRPC endpoint (e.g. http://localhost:4317). "
+            "When set, metrics/traces are pushed to an OTel collector. "
+            "When unset, falls back to Prometheus pull mode."
+        ),
+    )
+    group.add_argument(
+        "--event-bus-queue-size",
+        type=int,
+        default=10_000,
+        help=(
+            "Maximum number of events in the EventBus queue before "
+            "tail-drop. Default is 10000."
+        ),
+    )
+    group.add_argument(
+        "--prometheus-port",
+        type=int,
+        default=9090,
+        help=(
+            "Port for the Prometheus /metrics endpoint. "
+            "Only used when --otlp-endpoint is not set. Default is 9090."
+        ),
+    )
+    group.add_argument(
+        "--metrics-sample-rate",
+        type=float,
+        default=0.01,
+        help=(
+            "Fraction of chunks/blocks to track for lifecycle histograms "
+            "(0, 1.0]. Counters always count all events. Default is 0.01 (1%%)."
+        ),
+    )
+
+    # Lookup hash logging config
+    log_group = parser.add_argument_group(
+        "Lookup Hash Logging",
+        "Configuration for lookup hash file logging (offline analysis)",
+    )
+    log_group.add_argument(
+        "--lookup-hash-log-dir",
+        type=str,
+        default="",
+        help="Directory to write lookup hash JSONL files for offline analysis. "
+        "Empty string (default) disables logging.",
+    )
+    log_group.add_argument(
+        "--lookup-hash-log-rotation-interval",
+        type=int,
+        default=6 * 3600,
+        help="Time interval in seconds before rotating to a new log file. "
+        "Default is 21600 (6 hours).",
+    )
+    log_group.add_argument(
+        "--lookup-hash-log-rotation-max-size",
+        type=int,
+        default=100 * 1024 * 1024,
+        help="Max file size in bytes before rotating even if the time "
+        "interval has not elapsed. Default is 100MB (104857600).",
+    )
+    log_group.add_argument(
+        "--lookup-hash-log-max-files",
+        type=int,
+        default=100,
+        help="Max number of lookup hash log files to keep. "
+        "Oldest files are deleted when this limit is exceeded. Default is 100.",
+    )
+
+    trace_group = parser.add_argument_group(
+        "Trace Recording",
+        "Capture LMCache operations to a binary trace file for replay "
+        "(see `lmcache trace`).",
+    )
+    trace_group.add_argument(
+        "--trace-level",
+        type=str,
+        choices=["storage"],
+        default=None,
+        help="Enable trace recording at the given level. Currently only "
+        "'storage' is supported (records StorageManager public-API calls).",
+    )
+    trace_group.add_argument(
+        "--trace-output",
+        type=str,
+        default=None,
+        help="Path to write the trace file. Defaults to a timestamped "
+        "file under $TMPDIR when --trace-level is set without an explicit "
+        "output path.",
+    )
+
+    return parser
+
+
+def parse_args_to_observability_config(
+    args: argparse.Namespace,
+) -> ObservabilityConfig:
+    """Convert parsed command line arguments to an ObservabilityConfig.
+
+    Args:
+        args: Parsed arguments from the argument parser.
+
+    Returns:
+        The configuration object.
+    """
+    config = ObservabilityConfig(
+        enabled=not args.disable_observability,
+        max_queue_size=args.event_bus_queue_size,
+        metrics_enabled=not args.disable_metrics,
+        logging_enabled=not args.disable_logging,
+        tracing_enabled=args.enable_tracing,
+        otlp_endpoint=args.otlp_endpoint,
+        prometheus_port=args.prometheus_port,
+        metrics_sample_rate=args.metrics_sample_rate,
+        lookup_hash_log=LookupHashLogConfig(
+            output_dir=args.lookup_hash_log_dir,
+            rotation_interval_sec=args.lookup_hash_log_rotation_interval,
+            rotation_max_size=args.lookup_hash_log_rotation_max_size,
+            max_files=args.lookup_hash_log_max_files,
+        ),
+        trace_level=args.trace_level,
+        trace_output=args.trace_output,
+    )
+
+    if config.tracing_enabled and config.otlp_endpoint is None:
+        raise ValueError(
+            "--enable-tracing requires --otlp-endpoint to be set. "
+            "Tracing needs an OTLP gRPC endpoint to export spans."
+        )
+
+    return config
+
+
+def init_observability(obs_config: ObservabilityConfig) -> EventBus:
+    """Initialize OTel providers, EventBus, and register subscribers.
+
+    This is the single entry-point that every MP server calls at startup.
+    Returns a **started** EventBus.
+    """
+    # First Party
+    from lmcache.v1.mp_observability.event_bus import (
+        EventBusConfig,
+        init_event_bus,
+    )
+
+    # Set up OTel providers BEFORE creating subscribers so that
+    # module-level get_meter()/get_tracer() calls bind to the real provider
+    if obs_config.enabled and obs_config.metrics_enabled:
+        # First Party
+        from lmcache.v1.mp_observability.otel_init import init_otel_metrics
+
+        init_otel_metrics(
+            otlp_endpoint=obs_config.otlp_endpoint,
+            prometheus_port=obs_config.prometheus_port,
+        )
+
+    if obs_config.enabled and obs_config.tracing_enabled:
+        # First Party
+        from lmcache.v1.mp_observability.otel_init import init_otel_tracing
+
+        init_otel_tracing(otlp_endpoint=obs_config.otlp_endpoint)
+
+    bus = init_event_bus(
+        EventBusConfig(
+            enabled=obs_config.enabled,
+            max_queue_size=obs_config.max_queue_size,
+        )
+    )
+
+    if obs_config.metrics_enabled:
+        # First Party
+        from lmcache.v1.mp_observability.subscribers.metrics import (
+            L0LifecycleSubscriber,
+            L1LifecycleSubscriber,
+            L1MetricsSubscriber,
+            L2MetricsSubscriber,
+            SMMetricsSubscriber,
+        )
+
+        sample_rate = obs_config.metrics_sample_rate
+        bus.register_subscriber(L0LifecycleSubscriber(sample_rate=sample_rate))
+        bus.register_subscriber(L1MetricsSubscriber())
+        bus.register_subscriber(L1LifecycleSubscriber(sample_rate=sample_rate))
+        bus.register_subscriber(L2MetricsSubscriber())
+        bus.register_subscriber(SMMetricsSubscriber())
+
+    if obs_config.logging_enabled:
+        # First Party
+        from lmcache.v1.mp_observability.subscribers.logging import (
+            L1LoggingSubscriber,
+            L2LoggingSubscriber,
+            MPServerLoggingSubscriber,
+            SMLoggingSubscriber,
+        )
+
+        bus.register_subscriber(MPServerLoggingSubscriber())
+        bus.register_subscriber(L1LoggingSubscriber())
+        bus.register_subscriber(L2LoggingSubscriber())
+        bus.register_subscriber(SMLoggingSubscriber())
+
+    if obs_config.tracing_enabled:
+        # First Party
+        from lmcache.v1.mp_observability.subscribers.tracing import (
+            MPServerTracingSubscriber,
+            get_span_registry,
+        )
+
+        bus.register_subscriber(MPServerTracingSubscriber(get_span_registry()))
+
+    # Lookup hash file logging (independent of the logging_enabled flag —
+    # it has its own enable gate via output_dir).
+    if obs_config.lookup_hash_log.enabled:
+        # First Party
+        from lmcache.v1.mp_observability.subscribers.logging.lookup_hash import (
+            LookupHashLoggingSubscriber,
+        )
+
+        bus.register_subscriber(LookupHashLoggingSubscriber(obs_config.lookup_hash_log))
+
+    bus.start()
+    return bus