Files
2026-07-13 12:24:33 +08:00

290 lines
9.4 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""Per-qualname latency statistics for trace replay.
The replay driver times every dispatched call. Timings feed into this
collector, which computes count + mean + percentiles (p50, p90, p99)
per qualname.
The shape is deliberately simpler than
:class:`lmcache.cli.commands.bench.engine_bench.stats.StatsCollector`
— that one is tailored to OpenAI-style streaming inference (TTFT,
decode speed, etc.), which is not applicable to in-process storage
replay. Sharing the computation code between the two would push a
small helper deep into the bench module; keeping this collector local
keeps the storage-replay code self-contained and easy to evolve.
"""
# Future
from __future__ import annotations
# Standard
from dataclasses import dataclass, field
import csv
import json
import statistics
import threading
# First Party
from lmcache.logging import init_logger
logger = init_logger(__name__)
@dataclass
class OpStats:
"""Aggregate timing stats for one qualname.
Attributes:
qualname: The qualname these stats cover.
count: Number of successful replays for this qualname.
error_count: Number of replays that raised.
total_s: Total wall time spent replaying this qualname.
mean_ms: Mean per-call latency in milliseconds.
p50_ms: 50th-percentile latency in milliseconds.
p90_ms: 90th-percentile latency in milliseconds.
p99_ms: 99th-percentile latency in milliseconds.
min_ms: Minimum observed latency in milliseconds.
max_ms: Maximum observed latency in milliseconds.
"""
qualname: str
count: int
error_count: int
total_s: float
mean_ms: float
p50_ms: float
p90_ms: float
p99_ms: float
min_ms: float
max_ms: float
@dataclass
class _Bucket:
"""Internal per-qualname sample bucket."""
latencies_ms: list[float] = field(default_factory=list)
errors: int = 0
def _percentile(sorted_values: list[float], pct: float) -> float:
"""Return the *pct*-th percentile from an already-sorted list.
Uses nearest-rank with no interpolation: for N samples, the
p-th percentile is the sample at index ``ceil(p/100 * N) - 1``.
Returns 0.0 for empty input.
Args:
sorted_values: Ascending-sorted list of values.
pct: Percentile in ``[0, 100]``.
Returns:
The percentile value, or 0.0 for empty input.
"""
if not sorted_values:
return 0.0
if pct <= 0:
return sorted_values[0]
if pct >= 100:
return sorted_values[-1]
# bisect_left gives the insertion index; the nearest-rank formula
# maps p to ceil(p/100 * N) which equals floor((p/100 * N - eps) + 1).
n = len(sorted_values)
idx = max(0, min(n - 1, int((pct / 100.0) * n)))
return sorted_values[idx]
class ReplayStatsCollector:
"""Thread-safe per-qualname latency collector.
The replay driver is single-threaded at the dispatcher boundary,
but the underlying StorageManager performs async work on helper
threads whose timings may eventually feed back here. A lock keeps
concurrent ``record()`` calls safe.
"""
def __init__(self) -> None:
self._lock = threading.Lock()
self._buckets: dict[str, _Bucket] = {}
self._wall_start_s: float | None = None
self._wall_end_s: float | None = None
def mark_start(self, wall_time_s: float) -> None:
"""Record the wall-clock time when replay began.
Called once before the first ``record()``; replay-duration
metrics are derived from start/end marks.
Args:
wall_time_s: ``time.time()`` at replay start.
"""
with self._lock:
self._wall_start_s = wall_time_s
def mark_end(self, wall_time_s: float) -> None:
"""Record the wall-clock time when replay finished.
Args:
wall_time_s: ``time.time()`` at replay end.
"""
with self._lock:
self._wall_end_s = wall_time_s
def record(self, qualname: str, latency_s: float, failed: bool = False) -> None:
"""Record one replayed call.
Args:
qualname: Qualified name of the replayed function.
latency_s: Elapsed seconds for the call.
failed: ``True`` if the call raised. Failed calls still
contribute to the count but do not add a latency
sample — the raising path's timing is not comparable
to successful calls.
"""
with self._lock:
bucket = self._buckets.get(qualname)
if bucket is None:
bucket = _Bucket()
self._buckets[qualname] = bucket
if failed:
bucket.errors += 1
return
# Append is O(1); the one-shot sort in :meth:`summary` is
# O(N log N), which beats the O(N) shift from keeping the
# list sorted on insert. For large traces (>1M records
# per qualname) the driver should sample or switch to an
# approximation; for now, exact percentiles suffice.
bucket.latencies_ms.append(latency_s * 1000.0)
def total_duration_s(self) -> float:
"""Return replay wall-clock duration in seconds.
Returns:
``mark_end - mark_start`` if both were set, else 0.0.
"""
with self._lock:
if self._wall_start_s is None or self._wall_end_s is None:
return 0.0
return max(0.0, self._wall_end_s - self._wall_start_s)
def summary(self) -> dict[str, OpStats]:
"""Return a per-qualname :class:`OpStats` snapshot.
Returns:
A dict keyed by qualname. A qualname with only errors
still appears, with zero latency stats.
"""
with self._lock:
result: dict[str, OpStats] = {}
for qualname, bucket in self._buckets.items():
# Sort once per summary call — ``record`` keeps the
# list unsorted (O(1) append) so the total work is
# O(N log N) per summary rather than O(N) per insert.
lats = sorted(bucket.latencies_ms)
if lats:
mean = statistics.fmean(lats)
total_s = sum(lats) / 1000.0
p50 = _percentile(lats, 50)
p90 = _percentile(lats, 90)
p99 = _percentile(lats, 99)
lo, hi = lats[0], lats[-1]
else:
mean = total_s = p50 = p90 = p99 = lo = hi = 0.0
result[qualname] = OpStats(
qualname=qualname,
count=len(lats),
error_count=bucket.errors,
total_s=total_s,
mean_ms=mean,
p50_ms=p50,
p90_ms=p90,
p99_ms=p99,
min_ms=lo,
max_ms=hi,
)
return result
def export_csv(self, path: str) -> None:
"""Write per-qualname stats to a CSV file.
Columns: ``qualname, count, errors, mean_ms, p50_ms, p90_ms,
p99_ms, min_ms, max_ms``. One row per qualname.
Args:
path: File path to write. Overwritten if it exists.
"""
summary = self.summary()
with open(path, "w", newline="") as f:
w = csv.writer(f)
w.writerow(
[
"qualname",
"count",
"errors",
"mean_ms",
"p50_ms",
"p90_ms",
"p99_ms",
"min_ms",
"max_ms",
]
)
for qn in sorted(summary):
s = summary[qn]
w.writerow(
[
s.qualname,
s.count,
s.error_count,
f"{s.mean_ms:.6f}",
f"{s.p50_ms:.6f}",
f"{s.p90_ms:.6f}",
f"{s.p99_ms:.6f}",
f"{s.min_ms:.6f}",
f"{s.max_ms:.6f}",
]
)
def export_json(self, path: str) -> None:
"""Write per-qualname stats + replay duration to a JSON file.
Schema::
{
"duration_s": <float>,
"ops": {
"<qualname>": {
"count": int, "errors": int,
"mean_ms": float, "p50_ms": float,
"p90_ms": float, "p99_ms": float,
"min_ms": float, "max_ms": float
},
...
}
}
Args:
path: File path to write. Overwritten if it exists.
"""
summary = self.summary()
payload = {
"duration_s": self.total_duration_s(),
"ops": {
qn: {
"count": s.count,
"errors": s.error_count,
"mean_ms": s.mean_ms,
"p50_ms": s.p50_ms,
"p90_ms": s.p90_ms,
"p99_ms": s.p99_ms,
"min_ms": s.min_ms,
"max_ms": s.max_ms,
}
for qn, s in summary.items()
},
}
with open(path, "w") as f:
json.dump(payload, f, indent=2, sort_keys=True)
f.write("\n")