wehub-resource-sync
250 lines
8.5 KiB
Python
250 lines
8.5 KiB
Python
# 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
|