From 2e50567b8c32f349c0de81a4051eaef35cddfed9 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 9 Jul 2026 11:17:17 +0300 Subject: [PATCH] fix(proxy): flush billable-request counts on proxy shutdown PeriodicExportingMetricReader buffers up to one export interval of counts; without a final flush every restart silently dropped them. The factory registers the recorder it builds and proxy_shutdown_event pops and flushes it, bounded by a 5s timeout so a dead collector cannot stall shutdown --- .../enterprise_billing/billing_metrics.py | 41 ++++++++++++++++++- litellm/proxy/proxy_server.py | 10 +++++ .../test_billing_metrics.py | 30 +++++++++++++- 3 files changed, 79 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/enterprise_billing/billing_metrics.py b/litellm/proxy/enterprise_billing/billing_metrics.py index ea8c79509f6..4b2a7a22232 100644 --- a/litellm/proxy/enterprise_billing/billing_metrics.py +++ b/litellm/proxy/enterprise_billing/billing_metrics.py @@ -37,6 +37,7 @@ CLIENT_KEY_ENV = "LITELLM_BILLING_METRICS_CLIENT_KEY" CA_CERT_ENV = "LITELLM_BILLING_METRICS_CA_CERT" EXPORT_INTERVAL_ENV = "LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS" DEFAULT_EXPORT_INTERVAL_MS = 60_000 +SHUTDOWN_FLUSH_TIMEOUT_MS = 5_000 _METRICS_PATH = "/v1/metrics" METRIC_NAME = "litellm.enterprise.billable_requests" @@ -116,6 +117,11 @@ class BillingMetricsRecorder: def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: Optional[str]) -> None: self._counter.add(1, _billable_attributes(category, route, status_code, model_id)) + def shutdown(self) -> None: + """Final flush + exporter-thread stop. Without this, up to one export + interval of billable counts is dropped on every proxy restart.""" + self._provider.shutdown(timeout_millis=SHUTDOWN_FLUSH_TIMEOUT_MS) + def _export_interval_ms() -> int: raw = os.getenv(EXPORT_INTERVAL_ENV) @@ -176,6 +182,26 @@ def load_billing_metrics_config( ) +class _ActiveRecorderRegistry: + """One-slot registry linking the factory-built recorder to the shutdown + hook; the middleware instance holding the recorder is not reachable from + proxy_shutdown_event.""" + + def __init__(self) -> None: + self._recorder: Optional[BillingMetricsRecorder] = None + + def set(self, recorder: BillingMetricsRecorder) -> None: + self._recorder = recorder + + def pop(self) -> Optional[BillingMetricsRecorder]: + recorder = self._recorder + self._recorder = None + return recorder + + +_ACTIVE_RECORDER = _ActiveRecorderRegistry() + + def build_billing_metrics_recorder( *, premium: bool, license_data: Optional["EnterpriseLicenseData"], litellm_version: str ) -> Optional[BillingMetricsRecorder]: @@ -188,7 +214,20 @@ def build_billing_metrics_recorder( return None try: - return BillingMetricsRecorder(build_mtls_meter_provider(config)) + recorder = BillingMetricsRecorder(build_mtls_meter_provider(config)) except Exception as exc: # noqa: BLE001 -- metering must never break proxy startup verbose_proxy_logger.warning("Enterprise billing metrics disabled: failed to initialize exporter: %s", exc) return None + _ACTIVE_RECORDER.set(recorder) + return recorder + + +def shutdown_billing_metrics_recorder() -> None: + """Flush and stop the active recorder, if any. Idempotent; never raises.""" + recorder = _ACTIVE_RECORDER.pop() + if recorder is None: + return + try: + recorder.shutdown() + except Exception as exc: # noqa: BLE001 -- shutdown must never block or fail proxy exit + verbose_proxy_logger.warning("Enterprise billing metrics: final flush failed: %s", exc) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4e1bb0dcfc1..e228a152884 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -447,10 +447,15 @@ try: from litellm.proxy.enterprise_billing.billing_metrics import ( build_billing_metrics_recorder as _build_billing_metrics_recorder, ) + from litellm.proxy.enterprise_billing.billing_metrics import ( + shutdown_billing_metrics_recorder as _shutdown_billing_metrics_recorder, + ) build_billing_metrics_recorder: Optional[Callable[..., Optional[BillingRecorder]]] = _build_billing_metrics_recorder + shutdown_billing_metrics_recorder: Optional[Callable[[], None]] = _shutdown_billing_metrics_recorder except ImportError: build_billing_metrics_recorder = None + shutdown_billing_metrics_recorder = None from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -771,6 +776,11 @@ async def proxy_shutdown_event(): if db_writer_client is not None: await db_writer_client.close() # type: ignore[reportGeneralTypeIssues] + # final flush of billable-request counts: without it, up to one export + # interval of enterprise billing data is dropped on every restart + if shutdown_billing_metrics_recorder is not None: + shutdown_billing_metrics_recorder() + # flush remaining langfuse logs if "langfuse" in litellm.success_callback: try: diff --git a/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py b/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py index b89ee993742..66cd00f1fb2 100644 --- a/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py +++ b/tests/test_litellm/proxy/enterprise_billing/test_billing_metrics.py @@ -31,6 +31,8 @@ _ENV_VARS = ( def clear_env(monkeypatch): for name in _ENV_VARS: monkeypatch.delenv(name, raising=False) + yield + bm.shutdown_billing_metrics_recorder() def _write_certs(tmp_path: Path) -> Dict[str, str]: @@ -95,7 +97,33 @@ def test_premium_with_full_config_builds_recorder(monkeypatch, tmp_path): premium=True, license_data={"user_id": "org-1"}, litellm_version="1.0" ) assert isinstance(recorder, bm.BillingMetricsRecorder) - recorder._provider.shutdown() + bm.shutdown_billing_metrics_recorder() + + +def test_shutdown_flushes_active_recorder_once(monkeypatch, tmp_path): + """The shutdown hook must flush the recorder the factory built (buffered + counts are lost on restart otherwise) and be idempotent for repeat calls.""" + _set_full_env(monkeypatch, tmp_path) + shutdowns = [] + + class _SpyProvider: + def get_meter(self, name): + return MeterProvider().get_meter(name) + + def shutdown(self, timeout_millis=None): + shutdowns.append(timeout_millis) + + monkeypatch.setattr(bm, "build_mtls_meter_provider", lambda config: _SpyProvider()) + recorder = bm.build_billing_metrics_recorder(premium=True, license_data=None, litellm_version="1.0") + assert recorder is not None + + bm.shutdown_billing_metrics_recorder() + bm.shutdown_billing_metrics_recorder() + assert shutdowns == [bm.SHUTDOWN_FLUSH_TIMEOUT_MS] + + +def test_shutdown_without_active_recorder_is_noop(): + bm.shutdown_billing_metrics_recorder() # ── Config loading ────────────────────────────────────────────────────────────