diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 8a0edf4a6ca..5432c5095a2 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -46,6 +46,7 @@ - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"} - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"} - {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"} +- {id: quota_management.spend_tracking.surface_consistency.matches_every_surface, module: quota_management, tier: P1, behavior: spend_tracking, variant: surface_consistency, assertions: [matches_every_surface], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "One priced request lands the same response_cost on the spend log row, /key/info, /team/info, the usage export's /user/daily/activity/aggregated row, and the litellm_spend_metric Prometheus sample; each is a separate writer, so a rounding, dropped, or double-counted write on one drifts it from the rest (LIT-3620, LIT-5045)"} - {id: quota_management.spend_tracking.tags.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: tags, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Request tags round-trip to spend rows and tag rollups match tagged logs"} - {id: quota_management.spend_tracking.end_user.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: end_user, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "user= attribution lands the end-user id on the spend row"} - {id: quota_management.spend_tracking.per_model.writes_own_rows, module: quota_management, tier: P2, behavior: spend_tracking, variant: per_model, assertions: [writes_own_rows], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Each model on a shared key gets its own spend row"} diff --git a/tests/e2e/gateway/record_replay_ci_config.yml b/tests/e2e/gateway/record_replay_ci_config.yml index 5b3db0ff530..bc1f3f10562 100644 --- a/tests/e2e/gateway/record_replay_ci_config.yml +++ b/tests/e2e/gateway/record_replay_ci_config.yml @@ -2,3 +2,6 @@ general_settings: master_key: os.environ/LITELLM_MASTER_KEY store_model_in_db: true disable_model_info_refresh: true + +litellm_settings: + callbacks: ["prometheus"] diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index b7f59fe5f89..8b63b063e14 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -12,9 +12,10 @@ helpers from one place. from __future__ import annotations import time -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import Final from e2e_config import unique_marker @@ -48,6 +49,7 @@ from models import ( SpendLogsPageParams, SpendTagsResponse, TagSpend, + TeamInfoParams, UserDeleteBody, UserDeleteResponse, UserNewBody, @@ -57,6 +59,8 @@ from models import ( from proxy_client import Converged, ProxyClient, await_converged from pydantic import BaseModel, Field +METRICS_PATH: Final = "/metrics/" + __all__ = [ "BatchCreateBody", "CallbackLogMetadata", @@ -189,6 +193,7 @@ class DailyActivityKeyMetadata(BaseModel): class DailyActivityKeyMetrics(BaseModel): api_requests: int = 0 + spend: float = 0.0 class DailyActivityKeyBreakdown(BaseModel): @@ -209,6 +214,14 @@ class DailyActivityResponse(BaseModel): results: list[DailyActivityRow] = [] +class TeamInfoSpend(BaseModel): + spend: float | None = None + + +class TeamInfoSpendResponse(BaseModel): + team_info: TeamInfoSpend + + def _chat_body( model: str, content: str, @@ -334,6 +347,42 @@ class SpendClient: time.sleep(self.proxy.poll_interval) return spend + def team_spend(self, team_id: str) -> float: + return ( + unwrap( + self.proxy.transport.get( + "/team/info", + headers=self.proxy.transport.master, + params=TeamInfoParams(team_id=team_id), + response_type=TeamInfoSpendResponse, + ) + ).team_info.spend + or 0.0 + ) + + def poll_team_spend(self, team_id: str, *, minimum: float = 0.0) -> float: + outcome: Final = await_converged( + lambda: self.team_spend(team_id), + converged=lambda spend: spend > minimum, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + + def scrape_metrics(self) -> Mapping[str, ProbeResult]: + """GET /metrics/ on every replica in PROXY_REPLICA_URLS, keyed by replica. The + counter is per pod, so the union of the replicas is the fleet's exposition; the + trailing slash is the mounted app's own path, since bare /metrics answers a 307 + whose Location drops the port behind a Host-rewriting balancer.""" + return MappingProxyType( + { + replica: transport.probe(METRICS_PATH, params=NoBody()) + for replica, transport in self.proxy.replicas.items() + } + ) + def spend_logs_page( self, *, api_key: str | None, page: int, page_size: int ) -> SpendLogsPage: @@ -500,9 +549,21 @@ class SpendClient: return self.proxy.transport.probe("/health", params=HealthParams(model=model)) def daily_activity_for_key(self, token: str, *, start: datetime, end: datetime) -> DailyActivityKeyBreakdown | None: + return self._key_breakdown("/user/daily/activity", token, start=start, end=end) + + def usage_export_row_for_key( + self, token: str, *, start: datetime, end: datetime + ) -> DailyActivityKeyBreakdown | None: + """The key's row on /user/daily/activity/aggregated, the response the + dashboard's Export Usage Data CSV serializes.""" + return self._key_breakdown("/user/daily/activity/aggregated", token, start=start, end=end) + + def _key_breakdown( + self, route: str, token: str, *, start: datetime, end: datetime + ) -> DailyActivityKeyBreakdown | None: response: Final = unwrap( self.proxy.transport.get( - "/user/daily/activity", + route, headers=self.proxy.transport.master, params=DailyActivityParams( start_date=start.strftime("%Y-%m-%d"), @@ -519,9 +580,21 @@ class SpendClient: def poll_daily_activity_for_key( self, token: str, *, start: datetime, end: datetime, min_requests: int + ) -> DailyActivityKeyBreakdown | None: + return self._poll_key_breakdown(lambda: self.daily_activity_for_key(token, start=start, end=end), min_requests) + + def poll_usage_export_row_for_key( + self, token: str, *, start: datetime, end: datetime, min_requests: int + ) -> DailyActivityKeyBreakdown | None: + return self._poll_key_breakdown( + lambda: self.usage_export_row_for_key(token, start=start, end=end), min_requests + ) + + def _poll_key_breakdown( + self, fetch: Callable[[], DailyActivityKeyBreakdown | None], min_requests: int ) -> DailyActivityKeyBreakdown | None: outcome: Final = await_converged( - lambda: self.daily_activity_for_key(token, start=start, end=end), + fetch, converged=lambda found: found is not None and found.metrics.api_requests >= min_requests, timeout=self.proxy.poll_timeout, interval=self.proxy.poll_interval, diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_surface_consistency_e2e.py b/tests/e2e/quota_management/spend_tracking/test_spend_surface_consistency_e2e.py new file mode 100644 index 00000000000..9033b9d75c4 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_spend_surface_consistency_e2e.py @@ -0,0 +1,161 @@ +"""One priced request must land the same response_cost on every spend surface. + +A customer reconciles the bill from whichever surface they look at: the spend +log row, the key's and the team's rolled-up spend on /key/info and /team/info, +the usage page's Export Usage Data CSV (the dashboard serializes the +/user/daily/activity/aggregated rows it already holds; there is no server-side +CSV endpoint), and the litellm_spend_metric counter Prometheus scrapes. Each is +written by a different writer (the spend log insert, the key and team rollups in +db_spend_update_writer, the daily spend tables, the Prometheus success callback), +so one of them can drift without the others noticing: the cause of the +key-versus-log mismatch in LIT-3620 and the export-versus-console mismatch in +LIT-5045. The deployment carries its own per-token rates, so the expected cost +is computed from the returned usage rather than read off any one surface, and +every surface is held to that number. + +/metrics is per pod, so every replica the stack exports (PROXY_REPLICA_URLS) is +scraped directly and the samples merged; a stack that exports only its balancer +is scraped there until the pod that served the call answers. The request itself +is sent once. +""" + +from __future__ import annotations + +import time +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from itertools import groupby +from math import isclose +from types import MappingProxyType +from typing import Final + +import pytest +from e2e_config import provider_edge_base, unique_marker +from e2e_http import ProbeResult +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody +from prometheus_client.parser import text_string_to_metric_families +from proxy_client import Converged, await_converged +from spend_e2e_client import SpendClient, unwrap +from spend_reconciliation import INPUT_RATE, OUTPUT_RATE + +pytestmark = pytest.mark.e2e + +SPEND_METRIC: Final = "litellm_spend_metric_total" +KEY_HASH_LABEL: Final = "hashed_api_key" +TEAM_LABEL: Final = "team" + +SeriesLabels = tuple[tuple[str, str], ...] + + +def _spend_series_for_key(scrapes: Mapping[str, ProbeResult], token: str) -> Mapping[SeriesLabels, float]: + samples: Final = sorted( + (tuple(sorted(sample.labels.items())), sample.value) + for scrape in scrapes.values() + if scrape.status_code == 200 + for family in text_string_to_metric_families(scrape.body) + for sample in family.samples + if sample.name == SPEND_METRIC and sample.labels.get(KEY_HASH_LABEL) == token + ) + return MappingProxyType( + {labels: sum(value for _, value in group) for labels, group in groupby(samples, key=lambda sample: sample[0])} + ) + + +def _poll_spend_series_for_key(client: SpendClient, token: str) -> Mapping[SeriesLabels, float]: + outcome: Final = await_converged( + client.scrape_metrics, + converged=lambda scrapes: bool(_spend_series_for_key(scrapes, token)), + timeout=client.proxy.poll_timeout, + interval=client.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + scrapes: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result + assert _spend_series_for_key(scrapes, token), ( + f"{SPEND_METRIC} never exposed a series for {KEY_HASH_LABEL}={token} on any replica; " + f"last scrape status per replica: {({replica: scrape.status_code for replica, scrape in scrapes.items()})}" + ) + return _spend_series_for_key(scrapes, token) + + +def _same_spend(actual: float | None, expected: float) -> bool: + return actual is not None and isclose(actual, expected, rel_tol=1e-6, abs_tol=1e-9) + + +class TestSpendSurfaceConsistency: + @pytest.mark.replayable + @pytest.mark.covers("quota_management.spend_tracking.surface_consistency.matches_every_surface") + def test_one_request_lands_the_same_spend_on_every_surface( + self, client: SpendClient, resources: ResourceManager + ) -> None: + started: Final = datetime.now(timezone.utc) + marker: Final = unique_marker() + base: Final = provider_edge_base("openai") + model: Final = f"e2e-spend-surfaces-{marker}" + model_id: Final = client.proxy.create_model( + model, + LiteLLMParamsBody( + model="openai/gpt-5.6-luna", + api_key="os.environ/OPENAI_API_KEY", + api_base=None if base is None else f"{base}/v1", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-spend-surfaces-{marker}")) + resources.defer(lambda: client.proxy.delete_team(team_id)) + record: Final = client.generate_key_record( + KeyGenerateBody(team_id=team_id, models=[model], key_alias=f"e2e-spend-surfaces-{marker}") + ) + resources.defer(lambda: client.proxy.delete_key(record.key)) + assert record.token, "/key/generate answered without the key's token hash" + token: Final = record.token + + response: Final = unwrap( + client.proxy.chat( + record.key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"Reply with one word. {marker}")], + max_completion_tokens=128, + ), + ) + ) + usage: Final = response.usage + assert response.id, "successful response must have an ID" + assert usage is not None and usage.prompt_tokens and usage.completion_tokens, f"no billable usage: {usage}" + expected: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE + + rows: Final = client.proxy.poll_logs_for_request_id(response.id) + assert len(rows) == 1, f"expected one spend row for {response.id}, saw {len(rows)}: {rows}" + row: Final = rows[0] + assert row.api_key == token, f"spend row keyed by {row.api_key}, not the key's token hash {token}" + assert row.team_id == team_id, f"spend row attributed to team {row.team_id}, not {team_id}" + assert row.status == "success", f"spend row status {row.status}" + + key_spend: Final = client.poll_key_spend(record.key, minimum=expected * 0.999999) + team_spend: Final = client.poll_team_spend(team_id, minimum=expected * 0.999999) + export_row: Final = client.poll_usage_export_row_for_key( + token, start=started - timedelta(days=1), end=datetime.now(timezone.utc), min_requests=1 + ) + assert export_row is not None, f"/user/daily/activity/aggregated never listed key {token} under api_keys" + series: Final = _poll_spend_series_for_key(client, token) + off_team: Final = tuple(labels for labels in series if dict(labels).get(TEAM_LABEL) != team_id) + assert not off_team, f"{SPEND_METRIC} series for the key carry a team other than {team_id}: {off_team}" + + observed: Final = MappingProxyType( + { + "/spend/logs row": row.spend, + "/key/info spend": key_spend, + "/team/info spend": team_spend, + "usage export row (/user/daily/activity/aggregated)": export_row.metrics.spend, + SPEND_METRIC: sum(series.values()), + } + ) + drifted: Final = tuple(surface for surface, spend in observed.items() if not _same_spend(spend, expected)) + assert not drifted, ( + f"response_cost {expected} (usage {usage.prompt_tokens}x{INPUT_RATE} + " + f"{usage.completion_tokens}x{OUTPUT_RATE}) drifted on {drifted}; every surface: {dict(observed)}" + )