mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(savings): batch tracker persistence off the request hot path
record_request persisted the full savings state on every proxied request: json.dumps of up to 5000 history entries plus a blocking os.fsync, run synchronously under the shared metrics event loop. Concurrent sessions queued behind whichever request was mid-save. Serialize dominates that cost (~57%) and is GIL-bound, so offloading the write to a thread can't overlap it with the loop, and fsync-only batching caps at ~28%. Reducing save frequency is the only lever that helps. SavingsTracker gains save_flush_every (default 1, so direct and CLI callers still persist on every call). The proxy constructs it with 25 and flushes on graceful shutdown. Full-state writes make the counter a pure throttle: a skipped save loses nothing because the next write is a complete snapshot. Reads stay in-memory and never go stale; a hard crash loses at most 24 requests' lifetime delta on the proxy path. 5000-entry history, 1000 record_request calls: 6.10 -> 1.12 ms/call (5.4x), fsync syscalls 1000 -> 40 (25x).
This commit is contained in:
parent
e8151f059b
commit
ddfd662656
4 changed files with 99 additions and 3 deletions
|
|
@ -60,6 +60,12 @@ def _append_metric(
|
|||
)
|
||||
|
||||
|
||||
# The proxy persists savings state on every request. Batch that write so a busy
|
||||
# event loop isn't blocked re-serializing + fsyncing the whole history each time;
|
||||
# the tracker still flushes on graceful shutdown (see HeadroomProxy.shutdown).
|
||||
PROXY_SAVINGS_FLUSH_EVERY = 25
|
||||
|
||||
|
||||
class PrometheusMetrics:
|
||||
"""Prometheus-compatible metrics."""
|
||||
|
||||
|
|
@ -240,7 +246,9 @@ class PrometheusMetrics:
|
|||
|
||||
# Cumulative savings history (timestamp → cumulative tokens saved)
|
||||
self.savings_history: list[tuple[str, int]] = []
|
||||
self.savings_tracker = savings_tracker or SavingsTracker(stateless=stateless)
|
||||
self.savings_tracker = savings_tracker or SavingsTracker(
|
||||
stateless=stateless, save_flush_every=PROXY_SAVINGS_FLUSH_EVERY
|
||||
)
|
||||
self.cost_tracker = cost_tracker
|
||||
tracker_lifetime = self.savings_tracker.snapshot()["lifetime"]
|
||||
self._savings_tracker_input_tokens_offset = max(
|
||||
|
|
|
|||
|
|
@ -438,6 +438,7 @@ class SavingsTracker:
|
|||
max_response_history_points: int = DEFAULT_MAX_RESPONSE_HISTORY_POINTS,
|
||||
display_session_inactivity_minutes: int = (DEFAULT_DISPLAY_SESSION_INACTIVITY_MINUTES),
|
||||
stateless: bool = False,
|
||||
save_flush_every: int = 1,
|
||||
) -> None:
|
||||
# In stateless mode the tracker keeps live counters in memory but never
|
||||
# writes proxy_savings.json (honors HeadroomConfig.stateless, which
|
||||
|
|
@ -460,6 +461,13 @@ class SavingsTracker:
|
|||
),
|
||||
1,
|
||||
)
|
||||
# ponytail: per-record save throttle. Default 1 = persist every call
|
||||
# (the durable default that direct/CLI callers rely on). The async proxy
|
||||
# opts into a higher value so it doesn't json.dumps + fsync the whole
|
||||
# history on every request. Lossless because _save_locked always writes
|
||||
# the FULL state — a skipped save just means the next one is complete.
|
||||
self._save_flush_every = max(_coerce_int(save_flush_every, 1), 1)
|
||||
self._since_save = 0
|
||||
self._lock = threading.Lock()
|
||||
self._state = self._load_state()
|
||||
|
||||
|
|
@ -527,7 +535,7 @@ class SavingsTracker:
|
|||
}
|
||||
)
|
||||
self._trim_history_locked(reference_time=timestamp_dt)
|
||||
self._save_locked()
|
||||
self._maybe_save_locked()
|
||||
return True
|
||||
|
||||
def record_request(
|
||||
|
|
@ -661,7 +669,7 @@ class SavingsTracker:
|
|||
)
|
||||
self._trim_history_locked(reference_time=timestamp_dt)
|
||||
|
||||
self._save_locked()
|
||||
self._maybe_save_locked()
|
||||
return True
|
||||
|
||||
def _record_project_locked(
|
||||
|
|
@ -1003,7 +1011,27 @@ class SavingsTracker:
|
|||
|
||||
return compacted
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Persist any records held back by the save throttle.
|
||||
|
||||
Call on graceful shutdown so a batched proxy doesn't drop the tail of
|
||||
recent requests. No-op when nothing is buffered.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._since_save > 0:
|
||||
self._save_locked()
|
||||
|
||||
def _maybe_save_locked(self) -> None:
|
||||
"""Throttled persist: write only every ``_save_flush_every`` records.
|
||||
|
||||
Caller must hold ``self._lock``. Lossless by design — see ``__init__``.
|
||||
"""
|
||||
self._since_save += 1
|
||||
if self._since_save >= self._save_flush_every:
|
||||
self._save_locked()
|
||||
|
||||
def _save_locked(self) -> None:
|
||||
self._since_save = 0
|
||||
if self._stateless:
|
||||
# Stateless mode: live counters stay in memory; nothing is persisted.
|
||||
return
|
||||
|
|
|
|||
|
|
@ -1641,6 +1641,11 @@ class HeadroomProxy(
|
|||
# Stop all quota trackers via the registry
|
||||
await get_quota_registry().stop_all()
|
||||
|
||||
# Persist any savings the tracker's write throttle is still holding, so
|
||||
# a graceful shutdown doesn't drop the last few requests' totals.
|
||||
with contextlib.suppress(Exception):
|
||||
self.metrics.savings_tracker.flush()
|
||||
|
||||
# Print final stats
|
||||
self._print_summary()
|
||||
|
||||
|
|
|
|||
|
|
@ -1086,6 +1086,9 @@ def test_stats_history_persists_across_restarts_and_stats_stays_compatible(tmp_p
|
|||
assert full["history_summary"]["stored_points"] == 2
|
||||
assert full["history_summary"]["returned_points"] == 2
|
||||
|
||||
# The proxy batches savings writes, so force a flush before reading the
|
||||
# file directly mid-session (a graceful shutdown flushes automatically).
|
||||
client.app.state.proxy.metrics.savings_tracker.flush()
|
||||
persisted = json.loads(savings_path.read_text())
|
||||
assert persisted["lifetime"]["tokens_saved"] == 55
|
||||
assert persisted["lifetime"]["total_input_tokens"] == 240
|
||||
|
|
@ -1093,6 +1096,58 @@ def test_stats_history_persists_across_restarts_and_stats_stays_compatible(tmp_p
|
|||
assert persisted["display_session"]["requests"] == 2
|
||||
|
||||
|
||||
def test_savings_tracker_batches_saves_and_matches_immediate(tmp_path):
|
||||
"""save_flush_every batches disk writes; the threshold and flush() together
|
||||
produce the exact on-disk state an immediate (flush_every=1) tracker would.
|
||||
|
||||
Proves the batch boundary drops no data — the correctness half of the perf
|
||||
fix, independent of timing.
|
||||
"""
|
||||
events = [
|
||||
{
|
||||
"model": "gpt-4o",
|
||||
"input_tokens": 120,
|
||||
"tokens_saved": 10,
|
||||
"timestamp": "2026-03-27T09:00:00Z",
|
||||
},
|
||||
{
|
||||
"model": "gpt-4o",
|
||||
"input_tokens": 80,
|
||||
"tokens_saved": 5,
|
||||
"timestamp": "2026-03-27T09:01:00Z",
|
||||
},
|
||||
{
|
||||
"model": "gpt-4o",
|
||||
"input_tokens": 200,
|
||||
"tokens_saved": 25,
|
||||
"timestamp": "2026-03-27T09:02:00Z",
|
||||
},
|
||||
]
|
||||
|
||||
# Baseline: persists on every call (default save_flush_every=1).
|
||||
immediate_path = tmp_path / "immediate.json"
|
||||
immediate = SavingsTracker(path=str(immediate_path))
|
||||
for event in events:
|
||||
immediate.record_request(**event)
|
||||
|
||||
# Batched: writes only every 2 records; the tail lands on flush().
|
||||
batched_path = tmp_path / "batched.json"
|
||||
batched = SavingsTracker(path=str(batched_path), save_flush_every=2)
|
||||
|
||||
batched.record_request(**events[0])
|
||||
assert not batched_path.exists() # buffered, below threshold
|
||||
|
||||
batched.record_request(**events[1])
|
||||
assert batched_path.exists() # threshold reached, written
|
||||
|
||||
batched.record_request(**events[2]) # buffered again
|
||||
batched.flush() # tail persisted
|
||||
|
||||
assert json.loads(batched_path.read_text(encoding="utf-8")) == json.loads(
|
||||
immediate_path.read_text(encoding="utf-8")
|
||||
)
|
||||
|
||||
|
||||
def test_stats_history_csv_export_is_frontend_friendly(tmp_path, monkeypatch):
|
||||
savings_path = tmp_path / "proxy_savings.json"
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(savings_path))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue