# 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]