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

lmcache/v1/mp_observability/trace/codecs.py

@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Type codecs for trace argument serialization.
+
+The trace recorder needs to serialize arbitrary Python values that
+appear as arguments to decorated functions.  Msgpack natively handles
+``int``, ``float``, ``str``, ``bytes``, ``bool``, ``None``, ``list``,
+``tuple``, ``dict``.  Anything else needs an explicit codec.
+
+A codec is a pair ``(encode, decode)`` keyed on a Python type.  At
+encode time the value is wrapped in a ``{"__t__": tag, "v": payload}``
+dict so the decoder can recognize it without losing the round-trip.
+
+This registry is shared between the recorder (PR1, encode-only path
+exercised) and the replay driver (PR2, decode-only path).  Both halves
+ship together to keep the format and behavior coherent.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from dataclasses import dataclass
+from typing import Any, Callable
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey, PrefetchHandle
+
+
+@dataclass(frozen=True)
+class TypeCodec:
+    """Encode/decode pair for a single Python type."""
+
+    tag: str
+    encode: Callable[[Any], Any]
+    decode: Callable[[Any], Any]
+
+
+# Tag dispatch table populated by ``register_codec``.  Keyed by type.
+_BY_TYPE: dict[type, TypeCodec] = {}
+# Tag dispatch table for decode.  Keyed by tag string.
+_BY_TAG: dict[str, TypeCodec] = {}
+
+_WRAP_KEY = "__t__"
+_VALUE_KEY = "v"
+
+
+def register_codec(t: type, codec: TypeCodec) -> None:
+    """Register a codec for type ``t``.
+
+    Raises:
+        ValueError: If ``t`` or ``codec.tag`` is already registered.
+    """
+    if t in _BY_TYPE:
+        raise ValueError(f"codec already registered for type {t!r}")
+    if codec.tag in _BY_TAG:
+        raise ValueError(f"codec tag {codec.tag!r} already in use")
+    _BY_TYPE[t] = codec
+    _BY_TAG[codec.tag] = codec
+
+
+# ---------------------------------------------------------------------------
+# Encode / decode entry points
+# ---------------------------------------------------------------------------
+
+# Native msgpack types pass through unchanged.  Everything else must be
+# wrapped via a registered codec.
+_PASSTHROUGH = (int, float, str, bytes, bool, type(None))
+
+
+def encode_value(v: Any) -> Any:
+    """Encode ``v`` to a msgpack-friendly representation.
+
+    Recursively encodes lists, tuples, and dicts.  Tuples are preserved
+    via a tag so they can be decoded back to tuples (msgpack would
+    otherwise round-trip them as lists).
+
+    Raises:
+        TypeError: If ``v`` is of a type with no registered codec.
+    """
+    # Codec lookup by exact type takes priority so that registered
+    # types which happen to subclass ``tuple`` (e.g. ``torch.Size``) are
+    # handled by their codec rather than the generic tuple branch.
+    codec = _BY_TYPE.get(type(v))
+    if codec is not None:
+        return {_WRAP_KEY: codec.tag, _VALUE_KEY: codec.encode(v)}
+
+    if isinstance(v, _PASSTHROUGH):
+        return v
+    if isinstance(v, list):
+        return [encode_value(x) for x in v]
+    if isinstance(v, tuple):
+        return {_WRAP_KEY: "tuple", _VALUE_KEY: [encode_value(x) for x in v]}
+    if isinstance(v, dict):
+        # Dict keys must already be strings/ints for msgpack.  We do not
+        # encode keys, only values, to keep the on-wire form readable.
+        return {k: encode_value(x) for k, x in v.items()}
+
+    raise TypeError(
+        f"trace.codecs: no codec registered for type {type(v).__name__!r} "
+        f"(value={v!r}). Register one via register_codec() or extend "
+        f"the default registry."
+    )
+
+
+def decode_value(v: Any) -> Any:
+    """Decode a msgpack-deserialized value back to its native form.
+
+    Raises:
+        ValueError: If a wrapped value carries an unknown tag.
+    """
+    if isinstance(v, list):
+        return [decode_value(x) for x in v]
+    if isinstance(v, dict):
+        tag = v.get(_WRAP_KEY)
+        if tag is None:
+            return {k: decode_value(x) for k, x in v.items()}
+        if tag == "tuple":
+            return tuple(decode_value(x) for x in v[_VALUE_KEY])
+        codec = _BY_TAG.get(tag)
+        if codec is None:
+            raise ValueError(f"trace.codecs: unknown tag {tag!r}")
+        return codec.decode(v[_VALUE_KEY])
+    return v
+
+
+def encode_args(args: dict[str, Any]) -> dict[str, Any]:
+    """Encode an argument dict for serialization."""
+    return {k: encode_value(v) for k, v in args.items()}
+
+
+def decode_args(args: dict[str, Any]) -> dict[str, Any]:
+    """Decode an argument dict back to native values."""
+    return {k: decode_value(v) for k, v in args.items()}
+
+
+# ---------------------------------------------------------------------------
+# Default codecs for LMCache types
+# ---------------------------------------------------------------------------
+
+
+def _enc_object_key(k: ObjectKey) -> dict[str, Any]:
+    return {
+        "chunk_hash": k.chunk_hash,
+        "model_name": k.model_name,
+        "kv_rank": k.kv_rank,
+    }
+
+
+def _dec_object_key(d: dict[str, Any]) -> ObjectKey:
+    return ObjectKey(
+        chunk_hash=d["chunk_hash"],
+        model_name=d["model_name"],
+        kv_rank=d["kv_rank"],
+    )
+
+
+def _enc_layout_desc(d: MemoryLayoutDesc) -> dict[str, Any]:
+    return {
+        "shapes": [list(s) for s in d.shapes],
+        "dtypes": [str(dt) for dt in d.dtypes],
+    }
+
+
+# Mapping from str(torch.dtype) back to the dtype object.  Built lazily
+# the first time a layout desc is decoded.
+_DTYPE_BY_NAME: dict[str, torch.dtype] = {}
+
+
+def _resolve_dtype(name: str) -> torch.dtype:
+    if not _DTYPE_BY_NAME:
+        for attr in dir(torch):
+            obj = getattr(torch, attr)
+            if isinstance(obj, torch.dtype):
+                _DTYPE_BY_NAME[str(obj)] = obj
+    dtype = _DTYPE_BY_NAME.get(name)
+    if dtype is None:
+        raise ValueError(f"trace.codecs: unknown torch dtype {name!r}")
+    return dtype
+
+
+def _dec_layout_desc(d: dict[str, Any]) -> MemoryLayoutDesc:
+    return MemoryLayoutDesc(
+        shapes=[torch.Size(s) for s in d["shapes"]],
+        dtypes=[_resolve_dtype(dt) for dt in d["dtypes"]],
+    )
+
+
+def _enc_prefetch_handle(h: PrefetchHandle) -> dict[str, Any]:
+    return {
+        "prefetch_request_id": h.prefetch_request_id,
+        "external_request_id": h.external_request_id,
+        "l1_prefix_hit_count": h.l1_prefix_hit_count,
+        "total_requested_keys": h.total_requested_keys,
+        "submit_time": h.submit_time,
+    }
+
+
+def _dec_prefetch_handle(d: dict[str, Any]) -> PrefetchHandle:
+    return PrefetchHandle(
+        prefetch_request_id=d["prefetch_request_id"],
+        external_request_id=d["external_request_id"],
+        l1_prefix_hit_count=d["l1_prefix_hit_count"],
+        total_requested_keys=d["total_requested_keys"],
+        submit_time=d["submit_time"],
+    )
+
+
+def _enc_torch_size(s: torch.Size) -> list[int]:
+    return list(s)
+
+
+def _dec_torch_size(s: list[int]) -> torch.Size:
+    return torch.Size(s)
+
+
+def _enc_torch_dtype(dt: torch.dtype) -> str:
+    return str(dt)
+
+
+def _dec_torch_dtype(name: str) -> torch.dtype:
+    return _resolve_dtype(name)
+
+
+register_codec(
+    ObjectKey,
+    TypeCodec(tag="ObjectKey", encode=_enc_object_key, decode=_dec_object_key),
+)
+register_codec(
+    MemoryLayoutDesc,
+    TypeCodec(
+        tag="MemoryLayoutDesc",
+        encode=_enc_layout_desc,
+        decode=_dec_layout_desc,
+    ),
+)
+register_codec(
+    PrefetchHandle,
+    TypeCodec(
+        tag="PrefetchHandle",
+        encode=_enc_prefetch_handle,
+        decode=_dec_prefetch_handle,
+    ),
+)
+register_codec(
+    torch.Size,
+    TypeCodec(tag="torch.Size", encode=_enc_torch_size, decode=_dec_torch_size),
+)
+register_codec(
+    torch.dtype,
+    TypeCodec(tag="torch.dtype", encode=_enc_torch_dtype, decode=_dec_torch_dtype),
+)

lmcache/v1/mp_observability/trace/decorator.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""``@enable_tracing`` decorator for capturing function calls.
+
+The decorator publishes a single :data:`EventType.TRACE_CALL` event on
+**function entry** (inputs only).  Output values and exceptions are not
+captured — replay re-runs the function and observes the live outcome.
+
+The decorator imposes near-zero overhead when tracing is disabled: a
+single boolean attribute load is added to each call.  Argument
+introspection is performed only when the gate is on.
+
+Codecs that turn LMCache-specific argument types into msgpack-friendly
+forms live in :mod:`lmcache.v1.mp_observability.trace.codecs`.  The
+decorator deliberately does not import them; raw Python values are
+attached to the event and the recorder encodes at write time.  This
+keeps the decorator import-cheap and breaks an otherwise circular
+dependency (``StorageManager → decorator → codecs → StorageManager``).
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from functools import wraps
+from typing import Any, Callable, Sequence, TypeVar
+import inspect
+import time
+
+# First Party
+from lmcache.v1.mp_observability.event import Event, EventType
+from lmcache.v1.mp_observability.event_bus import get_event_bus
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+# Module-level gate.  Flipped on by the trace recorder when it
+# registers, off when it shuts down.  A simple bool is sufficient
+# (tracing capture is single-process; mutual visibility across threads
+# is not required for correctness — at worst a few events are missed
+# during the toggle window).
+_tracing_enabled: bool = False
+
+
+def is_tracing_enabled() -> bool:
+    """Return whether the trace gate is currently on."""
+    return _tracing_enabled
+
+
+def set_tracing_enabled(enabled: bool) -> None:
+    """Flip the trace gate.
+
+    Called by trace recorders during ``__init__`` (on) and ``close()``
+    (off).  Direct callers should not normally use this.
+    """
+    global _tracing_enabled
+    _tracing_enabled = enabled
+
+
+def publish_call_event(qualname: str, args: dict[str, Any]) -> None:
+    """Publish one ``TRACE_CALL`` event.
+
+    Used by :func:`enable_tracing` and by manual instrumentation
+    points (e.g. context-manager enter/exit) that cannot be wrapped
+    by the decorator.
+
+    Args:
+        qualname: Fully-qualified name of the call site.
+        args: Mapping of argument name to raw Python value.  Codec
+            encoding happens later, in the recorder.
+
+    ``time.monotonic()`` is sampled **here** (not on the drain thread)
+    so the recorded ``t_mono`` aligns with ``Event.timestamp`` —
+    otherwise the two clocks would drift by however long the drain
+    lagged behind the publisher.
+    """
+    if not _tracing_enabled:
+        return
+    t_mono = time.monotonic()
+    bus = get_event_bus()
+    bus.publish(
+        Event(
+            event_type=EventType.TRACE_CALL,
+            metadata={"qualname": qualname, "args": args, "t_mono": t_mono},
+        )
+    )
+
+
+def enable_tracing(
+    qualname: str | None = None,
+    capture: Sequence[str] | None = None,
+    redact: Sequence[str] = (),
+) -> Callable[[F], F]:
+    """Decorate a function so its calls publish ``TRACE_CALL`` events.
+
+    Args:
+        qualname: Fully-qualified call-site name placed in the event
+            metadata.  Defaults to ``f"{func.__module__}.{func.__qualname__}"``.
+        capture: If given, only these argument names are recorded.
+            ``None`` means capture every parameter except ``self`` and
+            ``cls``.
+        redact: Argument names that must not be recorded.  Applied
+            after ``capture``.
+
+    Returns:
+        A decorator that wraps the target function.
+
+    The signature is bound once at decoration time via
+    :func:`inspect.signature`, so per-call overhead is limited to a
+    bool check (when disabled) or a ``Signature.bind_partial`` plus
+    dict-comprehension (when enabled).
+    """
+    redact_set = frozenset(redact)
+    capture_set = frozenset(capture) if capture is not None else None
+
+    def deco(func: F) -> F:
+        sig = inspect.signature(func)
+        resolved_qualname = qualname or f"{func.__module__}.{func.__qualname__}"
+
+        # Pre-compute parameter names to record.  ``self`` and ``cls``
+        # are always dropped — recording the receiver object yields no
+        # useful information for replay and would force a codec for
+        # every receiver type.
+        param_names = [
+            name
+            for name in sig.parameters
+            if name not in ("self", "cls")
+            and (capture_set is None or name in capture_set)
+            and name not in redact_set
+        ]
+        param_set = frozenset(param_names)
+
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if _tracing_enabled:
+                bound = sig.bind_partial(*args, **kwargs)
+                bound.apply_defaults()
+                payload = {k: v for k, v in bound.arguments.items() if k in param_set}
+                publish_call_event(resolved_qualname, payload)
+            return func(*args, **kwargs)
+
+        # Expose the resolved metadata for tests and dispatcher
+        # registration.
+        wrapper.__lmc_trace_qualname__ = resolved_qualname  # type: ignore[attr-defined]
+        wrapper.__lmc_trace_params__ = tuple(param_names)  # type: ignore[attr-defined]
+        return wrapper  # type: ignore[return-value]
+
+    return deco

lmcache/v1/mp_observability/trace/format.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""On-disk format for trace files.
+
+A trace file is a length-prefixed msgpack stream:
+
+    [4-byte big-endian frame length][msgpack frame]
+    [4-byte big-endian frame length][msgpack frame]
+    ...
+
+The first frame is always a :class:`Header`.  All subsequent frames
+are :class:`Record` objects.  Length-prefixing keeps the reader
+simple and supports concurrent appenders (each frame is atomic on
+local filesystems for sizes below ``PIPE_BUF``).
+
+Format version is policed by the reader.  Unknown versions are
+rejected to make corrupt or future-format files fail loudly rather
+than silently producing garbage.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from typing import Any
+
+# Third Party
+import msgspec
+
+#: Magic bytes at the start of every file for sanity checking.  The
+#: reader rejects files that do not begin with these bytes.
+MAGIC: bytes = b"LMCT"
+
+#: Bumped whenever the on-wire framing layout changes in a backwards-
+#: incompatible way (length prefix, header/record struct shape, etc.).
+FORMAT_VERSION: int = 1
+
+#: Bumped whenever the captured API surface changes in a way that makes
+#: older traces undecodable or incorrect to replay — e.g. a traced
+#: StorageManager method gains/loses an argument, an argument type's
+#: codec wire form changes, or a new codec tag is introduced.  Owned by
+#: the trace subsystem, independent of the LMCache package version,
+#: because ``lmcache.__version__`` bumps cover many changes irrelevant
+#: to the trace contract.
+TRACE_SCHEMA_VERSION: int = 1
+
+
+class Header(msgspec.Struct, tag="header", omit_defaults=True):
+    """One-per-file metadata block."""
+
+    magic: bytes
+    """Always :data:`MAGIC`."""
+
+    format_version: int
+    """File format version; readers reject unknown values."""
+
+    level: str
+    """Trace level — currently ``"storage"``.  Future levels (``"mq"``,
+    ``"gpu"``) will share this format."""
+
+    trace_schema_version: int
+    """:data:`TRACE_SCHEMA_VERSION` at record time.  Replay drivers may
+    refuse mismatched schemas rather than silently misinterpreting old
+    traces."""
+
+    t_mono_start: float
+    """``time.monotonic()`` at recorder construction.  Record
+    timestamps are relative to this."""
+
+    t_wall_start: float
+    """``time.time()`` at recorder construction.  Used to correlate
+    with external logs / metrics in absolute wall-clock time."""
+
+    sm_config_json: str
+    """JSON dump of ``StorageManagerConfig`` at record time, or an
+    empty string when not available."""
+
+    sm_config_digest: str
+    """SHA-256 hex digest of :attr:`sm_config_json`.  Replay drivers
+    use this to detect mismatched configurations."""
+
+
+class Record(msgspec.Struct, tag="record", omit_defaults=True):
+    """One captured function call.
+
+    All records share the same shape; the ``qualname`` field
+    differentiates operations.  Future trace levels can introduce new
+    ``qualname`` values without bumping the format version.
+    """
+
+    t_mono: float
+    """Monotonic seconds since :attr:`Header.t_mono_start`."""
+
+    t_wall: float
+    """Wall-clock ``time.time()`` at the moment the event was
+    published."""
+
+    qualname: str
+    """Fully-qualified call-site name (e.g.
+    ``lmcache.v1.distributed.storage_manager.StorageManager.reserve_write``)."""
+
+    args: dict[str, Any]
+    """Codec-encoded argument dict.  See
+    :mod:`lmcache.v1.mp_observability.trace.codecs` for the codec
+    contract."""
+
+
+# msgspec encoders/decoders.  Reused per-process; both are thread-safe
+# after construction.
+_ENCODER = msgspec.msgpack.Encoder()
+_DECODER_HEADER = msgspec.msgpack.Decoder(Header)
+_DECODER_RECORD = msgspec.msgpack.Decoder(Record)
+
+
+def encode_header(h: Header) -> bytes:
+    """Serialize a header to msgpack bytes."""
+    return _ENCODER.encode(h)
+
+
+def encode_record(r: Record) -> bytes:
+    """Serialize a record to msgpack bytes."""
+    return _ENCODER.encode(r)
+
+
+def decode_header(buf: bytes) -> Header:
+    """Parse a header from msgpack bytes."""
+    return _DECODER_HEADER.decode(buf)
+
+
+def decode_record(buf: bytes) -> Record:
+    """Parse a record from msgpack bytes."""
+    return _DECODER_RECORD.decode(buf)

lmcache/v1/mp_observability/trace/lifecycle.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Trace recorder lifecycle helpers.
+
+Used by the cache server entry points (``server.py`` and
+``http_server.py``) to construct, register, and tear down trace
+recorders alongside the EventBus.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from datetime import datetime, timezone
+import os
+import tempfile
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.distributed.config import StorageManagerConfig
+from lmcache.v1.mp_observability.config import ObservabilityConfig
+from lmcache.v1.mp_observability.event_bus import EventBus
+from lmcache.v1.mp_observability.trace.recorder import (
+    StorageTraceRecorder,
+    TraceRecorder,
+)
+
+logger = init_logger(__name__)
+
+
+def _default_trace_path() -> str:
+    """Mint a timestamped path for an unnamed trace file."""
+    stamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
+    return os.path.join(
+        tempfile.gettempdir(),
+        f"lmcache-trace-{os.getpid()}-{stamp}.lct",
+    )
+
+
+def maybe_initialize_trace_recorder(
+    bus: EventBus,
+    obs_config: ObservabilityConfig,
+    storage_manager_config: StorageManagerConfig,
+) -> TraceRecorder | None:
+    """Construct and register a trace recorder if configured.
+
+    Args:
+        bus: The active EventBus to subscribe the recorder to.
+        obs_config: Observability config carrying the trace flags.
+        storage_manager_config: The StorageManagerConfig in use.  Used
+            to populate the trace file's header digest so a replay
+            driver can detect mismatched configurations.
+
+    Returns:
+        The created recorder, or ``None`` when ``obs_config.trace_level``
+        is unset.
+
+    The recorder is registered on the bus, so :meth:`EventBus.stop`
+    will invoke its ``shutdown`` (which flushes and closes the file).
+    Callers do not need to track the returned reference for cleanup;
+    it is returned only for testing and observation.
+    """
+    level = obs_config.trace_level
+    if not level:
+        return None
+    if level != "storage":
+        raise ValueError(
+            f"unsupported trace level {level!r}; only 'storage' is supported"
+        )
+
+    output_path = obs_config.trace_output or _default_trace_path()
+    if obs_config.trace_output is None:
+        logger.info(
+            "trace recording enabled (level=%s); no --trace-output given, "
+            "writing to %s",
+            level,
+            output_path,
+        )
+
+    recorder = StorageTraceRecorder(output_path=output_path)
+    recorder.attach_storage_config(storage_manager_config)
+    bus.register_subscriber(recorder)
+    return recorder