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:
Omar Gerardo 2026-07-05 12:51:55 +08:00
parent e8151f059b
commit ddfd662656
4 changed files with 99 additions and 3 deletions

View file

@ -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(

View file

@ -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

View file

@ -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()

View file

@ -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))