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

lmcache/cli/commands/trace/dispatch.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Dispatcher mapping recorded ``qualname`` strings to live callables.
+
+The recorder writes one :class:`~lmcache.v1.mp_observability.trace.format.Record`
+per decorated call, tagged by the function's fully-qualified name.
+The dispatcher translates those strings back into concrete calls on a
+live :class:`~lmcache.v1.distributed.storage_manager.StorageManager`.
+
+Adding support for a new traced operation is a two-line change: put
+``@enable_tracing`` on the function, then register a handler here with
+a matching ``qualname``.  No per-op schemas, no new event types.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from collections import deque
+from contextlib import AbstractContextManager
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.distributed.api import ObjectKey
+from lmcache.v1.distributed.storage_manager import StorageManager
+
+logger = init_logger(__name__)
+
+#: Fully-qualified name of the ``StorageManager`` class.  Used to
+#: build the qualnames of all its traced methods.  Kept as a constant
+#: so tests and dispatcher registrations stay in lock-step if the class
+#: is ever renamed or moved.
+_SM_PREFIX = "lmcache.v1.distributed.storage_manager.StorageManager"
+
+
+@dataclass
+class ReplayContext:
+    """State carried across dispatcher invocations during one replay.
+
+    The driver creates exactly one ``ReplayContext`` per trace file,
+    hands it to every dispatched handler, and closes the StorageManager
+    when replay finishes.
+
+    Attributes:
+        sm: The live StorageManager that receives replayed calls.
+        open_read_contexts: FIFO queue of
+            ``read_prefetched_results`` contexts entered but not yet
+            exited.  Matching ``__enter__``/``__exit__`` records are
+            popped in the order they were entered for the same
+            ``keys`` tuple.  A dict-of-deques keyed on
+            ``tuple(keys)`` supports interleaved contexts across
+            different key sets.
+    """
+
+    sm: StorageManager
+    open_read_contexts: dict[tuple[ObjectKey, ...], deque[AbstractContextManager]] = (
+        field(default_factory=dict)
+    )
+
+
+#: Type of a dispatcher handler: takes a :class:`ReplayContext` and an
+#: already-decoded ``args`` dict (keys = parameter names, values =
+#: native Python values restored by
+#: :mod:`lmcache.v1.mp_observability.trace.codecs`).  Handlers return
+#: nothing; any return value from the live call is discarded.
+Handler = Callable[[ReplayContext, dict[str, Any]], None]
+
+
+class CallDispatcher:
+    """Registry mapping recorded qualnames to replay handlers.
+
+    Handlers are plain callables — no per-op subclass hierarchy.  The
+    default factory :func:`build_default_dispatcher` populates the
+    registry with every v1 ``StorageManager`` operation.
+    """
+
+    def __init__(self) -> None:
+        self._handlers: dict[str, Handler] = {}
+
+    def register(self, qualname: str, handler: Handler) -> None:
+        """Register a handler for *qualname*.
+
+        Args:
+            qualname: Fully-qualified call-site name exactly matching
+                the ``qualname`` field written by the recorder.
+            handler: Callable invoked on each matching record.
+
+        Raises:
+            ValueError: If a handler is already registered for
+                *qualname*.
+        """
+        if qualname in self._handlers:
+            raise ValueError(f"handler already registered for {qualname!r}")
+        self._handlers[qualname] = handler
+
+    def has(self, qualname: str) -> bool:
+        """Return ``True`` if a handler is registered for *qualname*."""
+        return qualname in self._handlers
+
+    def registered_qualnames(self) -> list[str]:
+        """Return the list of currently registered qualnames.
+
+        Returns:
+            A new list of qualname strings, suitable for inspection or
+            assertion in tests.
+        """
+        return list(self._handlers)
+
+    def dispatch(
+        self,
+        qualname: str,
+        context: ReplayContext,
+        args: dict[str, Any],
+    ) -> None:
+        """Invoke the handler for *qualname*.
+
+        Args:
+            qualname: The recorded qualname.
+            context: Replay context passed through to the handler.
+            args: Decoded argument dict for the call.
+
+        Raises:
+            KeyError: If no handler is registered for *qualname*.  The
+                driver catches this and logs a warning so unknown
+                qualnames (e.g. from a future trace level) do not stop
+                replay.
+        """
+        handler = self._handlers.get(qualname)
+        if handler is None:
+            raise KeyError(qualname)
+        handler(context, args)
+
+
+# ---------------------------------------------------------------------------
+# Default handlers for v1 StorageManager operations
+# ---------------------------------------------------------------------------
+
+
+def _call_sm_method(method_name: str) -> Handler:
+    """Build a handler that forwards a record to ``sm.<method_name>``.
+
+    The returned callable invokes ``getattr(ctx.sm, method_name)(**args)``
+    and discards the result.  Used for every "plain" traced method on
+    StorageManager — ``reserve_write``, ``finish_write``,
+    ``submit_prefetch_task``, ``finish_read_prefetched``.
+
+    Args:
+        method_name: Attribute name on the live StorageManager.
+
+    Returns:
+        A :data:`Handler` closure.
+    """
+
+    def _handler(ctx: ReplayContext, args: dict[str, Any]) -> None:
+        method = getattr(ctx.sm, method_name)
+        method(**args)
+
+    _handler.__name__ = f"_call_sm_{method_name}"
+    return _handler
+
+
+def _enter_read_prefetched(ctx: ReplayContext, args: dict[str, Any]) -> None:
+    """Handle a ``read_prefetched_results.__enter__`` record.
+
+    Enters the live context manager and stashes it under
+    ``tuple(keys)``.  The matching ``__exit__`` handler pops the top
+    entry for that key tuple, preserving FIFO order when multiple
+    contexts are simultaneously open for identical key lists.
+
+    Args:
+        ctx: Active replay context.
+        args: Decoded record arguments.  Must contain ``"keys"``.
+    """
+    keys = args["keys"]
+    cm = ctx.sm.read_prefetched_results(keys)
+    cm.__enter__()
+    key_tuple = tuple(keys)
+    ctx.open_read_contexts.setdefault(key_tuple, deque()).append(cm)
+
+
+def _exit_read_prefetched(ctx: ReplayContext, args: dict[str, Any]) -> None:
+    """Handle a ``read_prefetched_results.__exit__`` record.
+
+    Pops and exits the oldest context opened for ``tuple(keys)``.  If
+    no matching open context exists — typically because the trace was
+    truncated between the enter and exit events — logs a warning and
+    continues so replay does not abort.
+
+    Args:
+        ctx: Active replay context.
+        args: Decoded record arguments.  Must contain ``"keys"``.
+    """
+    keys = args["keys"]
+    key_tuple = tuple(keys)
+    pending = ctx.open_read_contexts.get(key_tuple)
+    if not pending:
+        logger.warning(
+            "trace replay: read_prefetched_results.__exit__ with no "
+            "matching __enter__ (keys=%d); ignoring",
+            len(keys),
+        )
+        return
+    cm = pending.popleft()
+    if not pending:
+        del ctx.open_read_contexts[key_tuple]
+    # Pass a clean exit — replay does not reproduce caller-side
+    # exceptions.  The context manager's ``finally`` block runs
+    # regardless, so read locks are released.
+    cm.__exit__(None, None, None)
+
+
+def build_default_dispatcher() -> CallDispatcher:
+    """Return a :class:`CallDispatcher` populated with v1 handlers.
+
+    Covers every qualname the storage-level recorder emits:
+
+    * ``StorageManager.reserve_write``
+    * ``StorageManager.finish_write``
+    * ``StorageManager.submit_prefetch_task``
+    * ``StorageManager.finish_read_prefetched``
+    * ``StorageManager.read_prefetched_results.__enter__``
+    * ``StorageManager.read_prefetched_results.__exit__``
+
+    Returns:
+        A ready-to-use dispatcher.  Callers may further register
+        additional handlers on it for future trace levels.
+    """
+    dispatcher = CallDispatcher()
+    for method_name in (
+        "reserve_write",
+        "finish_write",
+        "submit_prefetch_task",
+        "finish_read_prefetched",
+    ):
+        dispatcher.register(
+            f"{_SM_PREFIX}.{method_name}",
+            _call_sm_method(method_name),
+        )
+    dispatcher.register(
+        f"{_SM_PREFIX}.read_prefetched_results.__enter__",
+        _enter_read_prefetched,
+    )
+    dispatcher.register(
+        f"{_SM_PREFIX}.read_prefetched_results.__exit__",
+        _exit_read_prefetched,
+    )
+    return dispatcher

lmcache/cli/commands/trace/driver.py

@@ -0,0 +1,372 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""High-level replay driver for storage-level traces.
+
+Usage::
+
+    driver = StorageReplayDriver(sm_config, trace_path)
+    result = driver.run()
+    driver.close()
+
+The driver:
+
+1. Opens the trace file via :class:`TraceReader`.
+2. Constructs a fresh :class:`StorageManager` from the supplied
+   :class:`StorageManagerConfig`.  Notably, the replay-side config is
+   chosen by the *caller*, not copied from the trace's header.  This
+   lets the same trace exercise different L1/L2 configurations —
+   e.g., record with a Redis L2 adapter and replay with a
+   local-filesystem adapter.
+3. Iterates records, decodes their argument dicts via the trace
+   codec registry, and dispatches each to a live StorageManager call
+   through a :class:`CallDispatcher`.  Each dispatch is aligned to
+   the recorded ``t_mono`` offset via ``time.sleep`` — the replay
+   never runs ahead of the recording.  Dispatching as-fast-as-
+   possible is unsafe because reads and writes are async in
+   ``StorageManager`` and carry cross-call dependencies; collapsing
+   the recorded gaps races the async queues.
+4. Records per-qualname timings into a :class:`ReplayStatsCollector`.
+
+Replay is deliberately single-threaded: the recorder captures calls
+in the order the EventBus drained them, which is already a
+linearization of the concurrent production calls.  Replaying
+in that same order preserves the observed interleaving without
+needing to reconstruct thread identities.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable
+import hashlib
+import json
+import time
+
+# First Party
+from lmcache.cli.commands.trace.dispatch import (
+    CallDispatcher,
+    ReplayContext,
+    build_default_dispatcher,
+)
+from lmcache.cli.commands.trace.stats import ReplayStatsCollector
+from lmcache.logging import init_logger
+from lmcache.v1.distributed.config import StorageManagerConfig
+from lmcache.v1.distributed.storage_manager import StorageManager
+from lmcache.v1.mp_observability.config import (
+    ObservabilityConfig,
+    init_observability,
+)
+from lmcache.v1.mp_observability.trace import codecs
+from lmcache.v1.mp_observability.trace.reader import TraceReader
+from lmcache.v1.mp_observability.trace.recorder import safe_storage_config_dict
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.mp_observability.event_bus import EventBus
+
+logger = init_logger(__name__)
+
+#: Default :class:`ObservabilityConfig` for replay sessions.
+#:
+#: Enables the EventBus and its logging subscribers so users can see
+#: SM/L1/L2 log output during replay.  Metrics and tracing are off:
+#: the OTel / Prometheus pipelines bind a port, which would collide
+#: across concurrent replays.  Declared at module scope to avoid
+#: Ruff B008 (mutable default argument).
+DEFAULT_REPLAY_OBS_CONFIG: ObservabilityConfig = ObservabilityConfig(
+    enabled=True,
+    metrics_enabled=False,
+    logging_enabled=True,
+    tracing_enabled=False,
+)
+
+
+@dataclass
+class ReplayResult:
+    """Summary returned from :meth:`StorageReplayDriver.run`.
+
+    Attributes:
+        records_replayed: Successful dispatches.
+        records_skipped: Records with no registered handler (likely
+            from a newer trace level).
+        records_failed: Records whose handler raised.
+        stats: The per-qualname timing collector.  Exposed so the
+            caller can export CSV/JSON or inspect individual
+            percentiles.
+        header_level: ``level`` field read from the trace header.
+        header_digest: ``sm_config_digest`` from the trace header.
+        replay_config_digest: SHA-256 of the replay-side
+            StorageManagerConfig, for mismatch comparisons.  Empty
+            string if the driver could not compute it.
+    """
+
+    records_replayed: int
+    records_skipped: int
+    records_failed: int
+    stats: ReplayStatsCollector
+    header_level: str
+    header_digest: str
+    replay_config_digest: str
+
+
+class StorageReplayDriver:
+    """Replays a storage-level trace against a live StorageManager.
+
+    The driver owns the StorageManager for its lifetime; call
+    :meth:`close` (or use as a context manager) to shut it down.
+    """
+
+    def __init__(
+        self,
+        sm_config: StorageManagerConfig,
+        trace_path: str,
+        dispatcher: CallDispatcher | None = None,
+        obs_config: ObservabilityConfig = DEFAULT_REPLAY_OBS_CONFIG,
+    ) -> None:
+        """Construct a driver.
+
+        Initializes the global observability EventBus **before**
+        constructing the StorageManager so internal events
+        (L0/L1/L2 lifecycle, eviction, etc.) flow through a live bus
+        during replay.  This lets the same logging and monitoring
+        subscribers that run in the real server attach to the replay
+        session — e.g. operators can eyeball L1/L2 log output to
+        spot eviction churn or L2 bottlenecks.
+
+        Metrics (OTel / Prometheus) are **off by default** because
+        the metrics pipeline binds a Prometheus port; running two
+        replays concurrently with the same port would fail.  Callers
+        who want metrics can pass their own ``obs_config``.
+
+        Args:
+            sm_config: Replay-side StorageManager configuration.
+                Determines L1 size, eviction policy, and L2 adapters
+                used during replay.  Typically different from the
+                recording-side config (the original deployment may
+                have used adapters unavailable on the replay host).
+            trace_path: Path to a ``.lct`` trace file written by
+                :class:`StorageTraceRecorder`.
+            dispatcher: Custom dispatcher.  When omitted, a fresh
+                :func:`build_default_dispatcher` is used; passing one
+                explicitly is useful for tests that register extra
+                handlers.
+            obs_config: Observability configuration for the replay
+                session.  Defaults to an enabled bus with logging
+                subscribers and no metrics/tracing.  The driver
+                installs this config as the global singleton via
+                :func:`init_observability` and stops the resulting
+                bus on :meth:`close`.
+        """
+        self._sm_config = sm_config
+        self._trace_path = trace_path
+        self._dispatcher = dispatcher or build_default_dispatcher()
+        self._closed = False
+
+        # Each resource is acquired under its own try/except so that a
+        # failure partway through __init__ still releases what has
+        # already been opened.  Without this, ``__exit__`` / ``close``
+        # never runs (the caller never got a valid instance), and the
+        # reader / bus would leak — see Cursor bugbot comment on
+        # PR #3075.
+        reader: TraceReader | None = None
+        bus: EventBus | None = None
+        try:
+            reader = TraceReader(trace_path)
+            bus = init_observability(obs_config)
+            self._sm = StorageManager(sm_config)
+        except BaseException:
+            # ``BaseException`` to also cover KeyboardInterrupt /
+            # SystemExit raised from inside StorageManager setup —
+            # the resources below must still be released.
+            if bus is not None:
+                try:
+                    bus.stop()
+                except Exception:
+                    logger.warning(
+                        "trace replay: error stopping bus during failed driver init",
+                        exc_info=True,
+                    )
+            if reader is not None:
+                try:
+                    reader.close()
+                except Exception:
+                    logger.warning(
+                        "trace replay: error closing reader during failed driver init",
+                        exc_info=True,
+                    )
+            raise
+
+        self._reader = reader
+        self._bus = bus
+
+    def __enter__(self) -> StorageReplayDriver:
+        return self
+
+    def __exit__(self, *_exc: object) -> None:
+        self.close()
+
+    # ------------------------------------------------------------------
+
+    @property
+    def trace_path(self) -> str:
+        """Path of the trace file being replayed."""
+        return self._trace_path
+
+    @property
+    def storage_manager(self) -> StorageManager:
+        """The live StorageManager driving replay.
+
+        Exposed primarily for tests that need to introspect residency
+        or eviction state after ``run()`` returns.
+        """
+        return self._sm
+
+    def close(self) -> None:
+        """Close the StorageManager, stop the bus, and close the reader.
+
+        Idempotent.  The order matters: the StorageManager may
+        publish teardown events, so the bus outlives it briefly.
+        """
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            self._sm.close()
+        finally:
+            try:
+                self._bus.stop()
+            finally:
+                self._reader.close()
+
+    # ------------------------------------------------------------------
+
+    def run(
+        self,
+        on_record: RecordCallback | None = None,
+    ) -> ReplayResult:
+        """Replay every record in the trace.
+
+        Dispatch is always paced to the recorded ``t_mono`` offsets
+        via ``time.sleep``: the replay never runs *ahead* of the
+        recording.  Running ahead is unsafe because ``StorageManager``
+        reads and writes are async — collapsing the recorded gaps
+        races the async queues and leads to retrieve misses.  If the
+        replay host is slower than recording, the loop simply lags
+        the recorded schedule.
+
+        Args:
+            on_record: Optional per-record callback invoked after
+                dispatch with ``(qualname, latency_s, failed)``.
+                Used by the CLI's ``--jsonl-out`` feature.
+
+        Returns:
+            A :class:`ReplayResult` summarizing the run.
+        """
+        stats = ReplayStatsCollector()
+        context = ReplayContext(sm=self._sm)
+        header = self._reader.header
+        t_start = time.time()
+        stats.mark_start(t_start)
+
+        replayed = skipped = failed = 0
+        t_wall_origin = time.monotonic()
+
+        for record in self._reader.records():
+            # Sleep just long enough to align to the recorded
+            # offset from the start of replay.  No speedup — if
+            # the replay machine is slower than recording, the
+            # loop simply runs behind.
+            target = t_wall_origin + record.t_mono
+            now = time.monotonic()
+            if now < target:
+                time.sleep(target - now)
+
+            try:
+                decoded_args = codecs.decode_args(record.args)
+            except Exception:
+                skipped += 1
+                logger.warning(
+                    "trace replay: failed to decode args for %s; skipping",
+                    record.qualname,
+                    exc_info=True,
+                )
+                if on_record is not None:
+                    on_record(record.qualname, 0.0, True)
+                continue
+
+            if not self._dispatcher.has(record.qualname):
+                skipped += 1
+                logger.warning(
+                    "trace replay: no handler for qualname %r; skipping",
+                    record.qualname,
+                )
+                if on_record is not None:
+                    on_record(record.qualname, 0.0, True)
+                continue
+
+            t0 = time.monotonic()
+            try:
+                self._dispatcher.dispatch(record.qualname, context, decoded_args)
+                latency = time.monotonic() - t0
+                stats.record(record.qualname, latency, failed=False)
+                replayed += 1
+                if on_record is not None:
+                    on_record(record.qualname, latency, False)
+            except Exception:
+                latency = time.monotonic() - t0
+                stats.record(record.qualname, latency, failed=True)
+                failed += 1
+                logger.warning(
+                    "trace replay: handler for %s raised",
+                    record.qualname,
+                    exc_info=True,
+                )
+                if on_record is not None:
+                    on_record(record.qualname, latency, True)
+
+        # Close any contexts the trace left open (truncated trace,
+        # or missing __exit__ records).  Releasing these keeps the
+        # StorageManager in a consistent state for a follow-up run
+        # or inspection.
+        for key_tuple, pending in list(context.open_read_contexts.items()):
+            while pending:
+                cm = pending.popleft()
+                try:
+                    cm.__exit__(None, None, None)
+                except Exception:
+                    logger.warning(
+                        "trace replay: forced exit of dangling "
+                        "read_prefetched_results raised (keys=%d)",
+                        len(key_tuple),
+                        exc_info=True,
+                    )
+            context.open_read_contexts.pop(key_tuple, None)
+
+        stats.mark_end(time.time())
+
+        # Digest of the replay-side config so callers can compare
+        # against ``header.sm_config_digest``.  Same hashing used by
+        # the recorder — see :func:`safe_storage_config_dict`.
+        safe = safe_storage_config_dict(self._sm_config)
+        replay_digest = hashlib.sha256(
+            json.dumps(safe, sort_keys=True).encode("utf-8")
+        ).hexdigest()
+
+        return ReplayResult(
+            records_replayed=replayed,
+            records_skipped=skipped,
+            records_failed=failed,
+            stats=stats,
+            header_level=header.level,
+            header_digest=header.sm_config_digest,
+            replay_config_digest=replay_digest,
+        )
+
+
+#: Callback signature for per-record hooks during replay.  Arguments
+#: are ``(qualname, latency_seconds, failed)``.  Declared at module
+#: scope so callers can type-annotate their hooks without importing
+#: the driver class.
+RecordCallback = Callable[[str, float, bool], None]