diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index bf593d404..c5ab25c64 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -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( diff --git a/headroom/proxy/savings_tracker.py b/headroom/proxy/savings_tracker.py index 30853ffad..926f9cd25 100644 --- a/headroom/proxy/savings_tracker.py +++ b/headroom/proxy/savings_tracker.py @@ -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,9 +1011,29 @@ 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: if self._stateless: # Stateless mode: live counters stay in memory; nothing is persisted. + self._since_save = 0 return try: self._path.parent.mkdir(parents=True, exist_ok=True) @@ -1035,6 +1063,9 @@ class SavingsTracker: except OSError: pass raise + # Reset only after a durable write. A failed save leaves the counter + # untouched so the next record retries instead of waiting a full window. + self._since_save = 0 except OSError as e: logger.warning("Failed to save savings history to %s: %s", self._path, e) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index a771360c9..3e1164001 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1649,6 +1649,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() diff --git a/tests/test_proxy_pipeline_lifecycle.py b/tests/test_proxy_pipeline_lifecycle.py index e68be2198..2a7a700e5 100644 --- a/tests/test_proxy_pipeline_lifecycle.py +++ b/tests/test_proxy_pipeline_lifecycle.py @@ -2,7 +2,7 @@ from __future__ import annotations import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock, call, patch +from unittest.mock import AsyncMock, Mock, call, patch import httpx from fastapi.testclient import TestClient @@ -92,6 +92,42 @@ def test_proxy_shutdown_unloads_image_models() -> None: quota_registry.stop_all.assert_awaited_once() +def test_proxy_shutdown_flushes_savings_tracker() -> None: + """Graceful shutdown must flush the savings tracker's batched tail. + + The proxy throttles savings persistence (save_flush_every=25), so buffered + requests only reach disk on the next threshold write or an explicit flush. + shutdown() is that flush; if the wiring regresses, a graceful stop silently + drops the last few requests' lifetime totals. The tracker's flush() logic is + covered in test_proxy_savings_history.py — this guards only the call site. + """ + config = ProxyConfig( + optimize=False, + image_optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + log_requests=False, + ccr_inject_tool=False, + ccr_handle_responses=False, + ccr_context_tracking=False, + ) + app = create_app(config) + proxy = app.state.proxy + proxy.http_client = None + proxy.memory_handler = None + proxy.metrics.savings_tracker.flush = Mock() + + quota_registry = SimpleNamespace(stop_all=AsyncMock()) + with ( + patch("headroom.proxy.server.get_quota_registry", return_value=quota_registry), + patch("headroom.models.ml_models.MLModelRegistry.unload_prefix"), + ): + asyncio.run(proxy.shutdown()) + + proxy.metrics.savings_tracker.flush.assert_called_once() + + def test_openai_chat_pipeline_events_cover_proxy_lifecycle(monkeypatch) -> None: recorder = _RecordingExtension() config = ProxyConfig( diff --git a/tests/test_proxy_savings_history.py b/tests/test_proxy_savings_history.py index b5b9ac048..ffc18d40d 100644 --- a/tests/test_proxy_savings_history.py +++ b/tests/test_proxy_savings_history.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import json import math +import tempfile from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace @@ -1086,6 +1087,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 +1097,89 @@ 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_failed_save_retries_on_next_record_not_after_full_window(tmp_path, monkeypatch): + """A transient write failure must not consume the flush window. + + The counter only resets after a durable write, so a save that raises leaves + it untouched and the next record retries immediately, rather than waiting + another save_flush_every calls. + """ + path = tmp_path / "proxy_savings.json" + tracker = SavingsTracker(path=str(path), save_flush_every=5) + + calls = {"n": 0} + real_mkstemp = tempfile.mkstemp + + def flaky_mkstemp(*args, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise OSError("simulated transient write failure") + return real_mkstemp(*args, **kwargs) + + monkeypatch.setattr(savings_tracker_module.tempfile, "mkstemp", flaky_mkstemp) + + for _ in range(5): + tracker.record_request(model="gpt-4o", input_tokens=10, tokens_saved=5) + assert not path.exists() # 5th call reached the threshold; its save failed + + # The 6th call must retry the save, not wait until the 10th. + tracker.record_request(model="gpt-4o", input_tokens=10, tokens_saved=5) + assert path.exists() + assert json.loads(path.read_text(encoding="utf-8"))["lifetime"]["requests"] == 6 + + 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))