From c8a2d8c3496ba643aa930b16982382d9a5d76d8c Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:00:36 +0000 Subject: [PATCH 1/8] feat(proxy): add LiteLLM_DailyGlobalSpend key-free rollup for the usage dashboard Adds a daily spend table without api_key or user_id, written atomically alongside LiteLLM_DailyUserSpend from the batched writer, reconciled from history by a scheduled job that advances a marker in LiteLLM_Config, and read by the key-free arm of the aggregated usage query once the marker covers the requested range. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 33 ++ .../litellm_proxy_extras/schema.prisma | 29 ++ litellm/constants.py | 3 + litellm/proxy/db/daily_spend_bulk_upsert.py | 98 +++-- litellm/proxy/db/db_spend_update_writer.py | 7 +- .../common_daily_activity.py | 38 +- litellm/proxy/proxy_server.py | 43 ++ litellm/proxy/schema.prisma | 29 ++ .../daily_global_spend_rollup.py | 235 +++++++++++ schema.prisma | 29 ++ .../proxy/db/test_daily_spend_bulk_upsert.py | 150 +++++++ .../proxy/db/test_db_spend_update_writer.py | 101 ++++- .../test_common_daily_activity.py | 154 ++++++- .../proxy/proxy_server/test_lifecycle.py | 48 +++ .../test_daily_global_spend_rollup.py | 382 ++++++++++++++++++ 15 files changed, 1347 insertions(+), 32 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql create mode 100644 litellm/proxy/spend_tracking/daily_global_spend_rollup.py create mode 100644 tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql new file mode 100644 index 00000000000..1d6cdea0c7b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql @@ -0,0 +1,33 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGlobalSpend" ( + "id" TEXT NOT NULL, + "date" TEXT NOT NULL, + "model" TEXT, + "model_group" TEXT, + "custom_llm_provider" TEXT, + "mcp_namespaced_tool_name" TEXT, + "endpoint" TEXT, + "prompt_tokens" BIGINT NOT NULL DEFAULT 0, + "completion_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "api_requests" BIGINT NOT NULL DEFAULT 0, + "successful_requests" BIGINT NOT NULL DEFAULT 0, + "failed_requests" BIGINT NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyGlobalSpend_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_idx" ON "LiteLLM_DailyGlobalSpend"("date"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_model_model_group_custom_llm__key" ON "LiteLLM_DailyGlobalSpend"("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8072df5aa5b..5d433e916d6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -809,6 +809,35 @@ model LiteLLM_DailyUserSpend { @@index([endpoint]) } +// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view +model LiteLLM_DailyGlobalSpend { + id String @id @default(uuid()) + date String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@index([date]) +} + // Track daily organization spend metrics per model and key model LiteLLM_DailyOrganizationSpend { id String @id @default(uuid()) diff --git a/litellm/constants.py b/litellm/constants.py index 565c6433c6e..1bf3a150aeb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2034,6 +2034,9 @@ PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90 # Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide # expiry cannot produce an alert too large for the channel delivering it. PTU_LAPSED_ALERT_LIMIT: Final[int] = 10 +DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID: Final[str] = "daily_global_spend_reconcile_job" +DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS: Final[int] = 3600 +DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM: Final[str] = "daily_global_spend_reconciled_through" # Slack allowed when deciding a sentinel row is stale. The row's updated_at and the # run's cutoff are stamped by different hosts, so clock skew between them must not let # one run delete a charge another just wrote. A stale row is hours old and a concurrent diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index a143643577e..c83043101eb 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -25,29 +25,41 @@ SpendRow = Mapping[str, object] @dataclass(frozen=True, slots=True) class DailySpendTable: - """The physical table behind one entity's daily rollup.""" + """A daily rollup table and the unique constraint its upserts arbitrate on.""" name: str - entity_id_column: str + key_columns: tuple[str, ...] carries_request_id: bool = False -DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( - { - "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), - "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), - "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), - "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), - "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), - "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), - } -) - # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. _KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") + +def _entity_table(name: str, entity_id_column: str, carries_request_id: bool = False) -> DailySpendTable: + return DailySpendTable( + name=name, key_columns=(entity_id_column, *_KEY_COLUMNS), carries_request_id=carries_request_id + ) + + +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": _entity_table("LiteLLM_DailyUserSpend", "user_id"), + "team": _entity_table("LiteLLM_DailyTeamSpend", "team_id"), + "org": _entity_table("LiteLLM_DailyOrganizationSpend", "organization_id"), + "end_user": _entity_table("LiteLLM_DailyEndUserSpend", "end_user_id"), + "agent": _entity_table("LiteLLM_DailyAgentSpend", "agent_id"), + "tag": _entity_table("LiteLLM_DailyTagSpend", "tag", carries_request_id=True), + } +) + +GLOBAL_SPEND_TABLE: Final = DailySpendTable( + name="LiteLLM_DailyGlobalSpend", + key_columns=("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"), +) + _COUNTER_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -92,7 +104,7 @@ def _as_float(value: object) -> float: def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) + return tuple(_as_text(transaction.get(column)) for column in table.key_columns) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -130,7 +142,11 @@ def _row_params( return ( str(uuid.uuid4()), *key, - None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), + *( + () + if "model_group" in table.key_columns + else (None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),) + ), *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), @@ -140,26 +156,25 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - table.entity_id_column, - *_KEY_COLUMNS, - "model_group", + *table.key_columns, + *(() if "model_group" in table.key_columns else ("model_group",)), *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), ) -def build_bulk_upsert( +def _upsert_statement( table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """The single statement writing one merged batch, plus its positional arguments.""" + first_param: int, +) -> str: columns: Final = _insert_columns(table) quoted_table: Final = f'"{table.name}"' rows: Final = ", ".join( "(" + ", ".join( - f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" + f"${first_param + row_index * len(columns) + offset}::{_CASTS.get(column, 'text')}" for offset, column in enumerate(columns) ) + ", (NOW() AT TIME ZONE 'UTC'))" @@ -176,11 +191,44 @@ def build_bulk_upsert( if table.carries_request_id else "" ) - sql: Final = ( + return ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted(table.key_columns)}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) - return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) + + +def _params(table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]]) -> tuple[SqlValue, ...]: + return tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) + + +def build_bulk_upsert( + table: DailySpendTable, + batch: Sequence[tuple[tuple[str, ...], SpendRow]], +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" + return _upsert_statement(table, batch, first_param=1), _params(table, batch) + + +def build_bulk_upsert_with_global_rollup( + table: DailySpendTable, + batch: Sequence[tuple[tuple[str, ...], SpendRow]], +) -> tuple[str, tuple[SqlValue, ...]]: + """One statement writing a batch to its table and, atomically, its key-free rollup + to ``LiteLLM_DailyGlobalSpend``. + + A data-modifying CTE runs both inserts in the same snapshot and transaction, so a + batch that lands in one table lands in both and a retried deadlock replays both. + Postgres does not order the CTE against the main statement, so two writers can still + deadlock across the tables; the caller's deadlock retry covers that, and each insert + takes its own rows in key order so same-table lock order stays deterministic. + """ + global_batch: Final = merge_by_conflict_key(GLOBAL_SPEND_TABLE, tuple(row for _, row in batch)) + entity_params: Final = _params(table, batch) + sql: Final = ( + f"WITH entity_rows AS (\n{_upsert_statement(table, batch, first_param=1)}\nRETURNING 1)\n" + f"{_upsert_statement(GLOBAL_SPEND_TABLE, global_batch, first_param=len(entity_params) + 1)}" + ) + return sql, (*entity_params, *_params(GLOBAL_SPEND_TABLE, global_batch)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index eaa03c5d7f7..d5c839a9be8 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -46,6 +46,7 @@ from litellm.proxy._types import ( from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, + build_bulk_upsert_with_global_rollup, merge_by_conflict_key, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -1939,7 +1940,11 @@ class DBSpendUpdateWriter: merged_batch = merge_by_conflict_key( table=table, transactions=tuple(transactions_to_process.values()) ) - sql, params = build_bulk_upsert(table=table, batch=merged_batch) + sql, params = ( + build_bulk_upsert_with_global_rollup(table=table, batch=merged_batch) + if entity_type == "user" + else build_bulk_upsert(table=table, batch=merged_batch) + ) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 8a3ba196ab2..f1d78dca201 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -11,6 +11,8 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE +from litellm.proxy.spend_tracking.daily_global_spend_rollup import reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, @@ -734,6 +736,30 @@ def _rollup_metric_select(table_name: str) -> str: _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" +async def key_free_source_table(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: + """The table the key-free arm reads from, when the global rollup can answer instead of the per-key table. + + Only an unfiltered read of the user table has the same rows as ``LiteLLM_DailyGlobalSpend``, + and only through the day the reconcile marker has reached: the writer keeps that day + current, later days are covered once the next run advances the marker. + """ + if query["table_name"] != "litellm_dailyuserspend": + return None + if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: + return None + _, adjusted_end = _adjust_dates_for_timezone( + query["start_date"], query["end_date"], query["timezone_offset_minutes"], query["include_current_utc_day"] + ) + try: + marker: Final = await reconciled_through(prisma_client) + except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read + verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) + return None + if marker is None or adjusted_end > marker: + return None + return GLOBAL_SPEND_TABLE.name + + def _build_aggregated_sql_query( *, table_name: str, @@ -746,13 +772,16 @@ def _build_aggregated_sql_query( exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path timezone_offset_minutes: int | None = None, include_current_utc_day: bool = False, + key_free_table: str | None = None, ) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params """Build the GROUPING SETS query for aggregated daily activity. One statement, two UNION ALL arms over the same WHERE clause. The first arm is key-free: grand total, per-date totals and the (date, model / model_group / provider / mcp / endpoint) rollups, so its row count never grows with the number - of keys. The second arm emits the (date, , api_key) rollups for the + of keys; it reads ``key_free_table`` when given (the global rollup, whose row count + never grew with the number of keys to begin with) and the entity table otherwise. + The second arm emits the (date, , api_key) rollups for the USAGE_TOP_API_KEYS_LIMIT highest-spend keys only. Both arms share the 7-bit group_level bitmask (date, api_key, model, model_group, provider, mcp, endpoint). @@ -778,6 +807,7 @@ def _build_aggregated_sql_query( ) sentinel_param: Final = f"${len(where_params) + 1}" metric_select: Final = _rollup_metric_select(table_name) + key_free_source: Final = key_free_table or pg_table # TODO: drop the successful_requests/failed_requests aggregates (and the # total_successful_requests metadata they feed) once the admin UI reads SGR @@ -796,7 +826,7 @@ def _build_aggregated_sql_query( | GROUPING(model, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level,{metric_select} - FROM "{pg_table}" + FROM "{key_free_source}" WHERE {where_clause} GROUP BY GROUPING SETS ( (date), @@ -1387,7 +1417,9 @@ async def get_daily_activity_aggregated( timezone_offset_minutes=timezone_offset_minutes, include_current_utc_day=include_current_utc_day, ) - sql_query, sql_params = _build_aggregated_sql_query(**query_kwargs) + sql_query, sql_params = _build_aggregated_sql_query( + **query_kwargs, key_free_table=await key_free_source_table(prisma_client, query_kwargs) + ) entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None raw_rows, raw_entity_rows = await asyncio.gather( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f63e088ebf7..ed8f6886734 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -259,6 +259,7 @@ from litellm.constants import ( APSCHEDULER_MISFIRE_GRACE_TIME, APSCHEDULER_REPLACE_EXISTING, CLI_SSO_SESSION_TTL_SECONDS, + DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, DAYS_IN_A_MONTH, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_MODEL_CREATED_AT_TIME, @@ -662,6 +663,9 @@ from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + run_scheduled_daily_global_spend_reconcile, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, active_spend_counter_batch, @@ -9970,6 +9974,12 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) + cls._initialize_daily_global_spend_reconcile_job( + scheduler=scheduler, + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma_client, + ) + ### PTU DAILY ROLLUP ### from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, @@ -10311,6 +10321,39 @@ class ProxyStartupEvent: "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" ) + @classmethod + def _initialize_daily_global_spend_reconcile_job( + cls, + scheduler: AsyncIOScheduler, + proxy_logging_obj: ProxyLogging, + prisma_client: PrismaClient, + ) -> None: + async def alert(message: str) -> None: + await proxy_logging_obj.alerting_handler( + message=message, + level="High", + alert_type=AlertType.failed_tracking_spend, + ) + + async def reconcile() -> None: + await run_scheduled_daily_global_spend_reconcile( + prisma_client, + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, + alert=alert, + ) + + scheduler.add_job( + reconcile, + "cron", + hour=0, + minute=30, + timezone="UTC", + id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + next_run_time=datetime.now(timezone.utc) + timedelta(minutes=2), + ) + @classmethod async def _initialize_slack_alerting_jobs( cls, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8072df5aa5b..5d433e916d6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -809,6 +809,35 @@ model LiteLLM_DailyUserSpend { @@index([endpoint]) } +// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view +model LiteLLM_DailyGlobalSpend { + id String @id @default(uuid()) + date String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@index([date]) +} + // Track daily organization spend metrics per model and key model LiteLLM_DailyOrganizationSpend { id String @id @default(uuid()) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py new file mode 100644 index 00000000000..9d344421332 --- /dev/null +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -0,0 +1,235 @@ +"""Reconcile ``LiteLLM_DailyGlobalSpend`` from ``LiteLLM_DailyUserSpend``, one day per transaction. + +The spend writer keeps both tables in step from the moment it is deployed; this job rolls up +the days before that and records how far it has reached in ``LiteLLM_Config`` so usage reads +know when the global table can answer for a date range. It runs as a background cron, never +in a Prisma migration, since on a large deployment the aggregate is minutes of work. +""" + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import date, datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, + DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, +) +from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE +from litellm.repositories.config_repository import ConfigRepository + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from litellm.proxy.utils import PrismaClient + +_DAY_TRANSACTION_TIMEOUT: Final = timedelta(minutes=10) +_REPLAY_DAYS: Final = 1 +_METRIC_COLUMNS: Final = ( + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", + "api_requests", + "successful_requests", + "failed_requests", + "compression_savings_spend", + "prompt_caching_savings_spend", + "gateway_injected_caching_savings_spend", + "autorouter_savings_spend", + "spend", +) + + +def _quoted(columns: tuple[str, ...]) -> str: + return ", ".join(f'"{column}"' for column in columns) + + +def _reconcile_day_sql() -> str: + key_columns: Final = GLOBAL_SPEND_TABLE.key_columns + normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in key_columns) + sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) + overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) + return ( + f'INSERT INTO "{GLOBAL_SPEND_TABLE.name}" ("id", {_quoted(key_columns)}, {_quoted(_METRIC_COLUMNS)}, ' + '"updated_at")\n' + f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" + 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' + f"GROUP BY {normalized_keys}\n" + f"ON CONFLICT ({_quoted(key_columns)}) DO UPDATE SET {overwrite}, " + "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" + ) + + +RECONCILE_DAY_SQL: Final = _reconcile_day_sql() +_LOCK_GLOBAL_TABLE_SQL: Final = f'LOCK TABLE "{GLOBAL_SPEND_TABLE.name}" IN EXCLUSIVE MODE' +_PENDING_DAYS_SQL: Final = ( + 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' +) + + +class ReconciledThrough(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + reconciled_through: str + + +class _MarkerRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + param_value: object = None + + +class _DateRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + date: str + + +@dataclass(frozen=True, slots=True) +class ReconcileResult: + days_reconciled: tuple[str, ...] + reconciled_through: str | None + failed_day: str | None = None + + +def _marker_from_param_value(value: object) -> str | None: + try: + parsed: Final = ( + ReconciledThrough.model_validate_json(value) + if isinstance(value, str) + else ReconciledThrough.model_validate(value) + ) + except ValidationError: + return None + return parsed.reconciled_through + + +async def reconciled_through(prisma_client: "PrismaClient") -> str | None: + """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" + from litellm.proxy.utils import get_config_param + + row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) + + +async def _record_reconciled_through(prisma_client: "PrismaClient", day: str) -> None: + from litellm.proxy.utils import invalidate_config_param + + await ConfigRepository(prisma_client).set_param( + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ReconciledThrough(reconciled_through=day).model_dump_json() + ) + await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +def _first_pending_day(marker: str | None) -> str: + if marker is None: + return "" + return (date.fromisoformat(marker) - timedelta(days=_REPLAY_DAYS)).isoformat() + + +async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: + """Every UTC day through today still to roll up, oldest first; the marker day and the one + before it are replayed so rows flushed by a pre-writer pod during a rolling deploy are folded in.""" + marker: Final = await reconciled_through(prisma_client) + rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), today.isoformat()) + return tuple(sorted({*(_DateRow.model_validate(row).date for row in rows), today.isoformat()})) + + +async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: + """Rewrite one day of the global table from the per-key sums; the table lock keeps the + writer's increments out between the aggregate and the overwrite so none are lost.""" + async with prisma_client.db.tx(timeout=_DAY_TRANSACTION_TIMEOUT) as transaction: + await transaction.execute_raw(_LOCK_GLOBAL_TABLE_SQL) + await transaction.execute_raw(RECONCILE_DAY_SQL, day) + + +async def run_daily_global_spend_reconcile( + prisma_client: "PrismaClient", + today: date | None = None, +) -> ReconcileResult: + """Roll up every pending day, advancing the marker after each; a failing day stops the run + with the marker on the last good day so the next run resumes there.""" + effective_today: Final = today or datetime.now(timezone.utc).date() + days: Final = await pending_days(prisma_client, effective_today) + done: Final = await _reconcile_until_failure(prisma_client, days) + failed: Final = days[len(done)] if len(done) < len(days) else None + marker: Final = done[-1] if done else await reconciled_through(prisma_client) + return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=failed) + + +async def _reconcile_until_failure(prisma_client: "PrismaClient", days: tuple[str, ...]) -> tuple[str, ...]: + for index, day in enumerate(days): + if not await _reconcile_and_record(prisma_client, day): + return days[:index] + return days + + +async def _reconcile_and_record(prisma_client: "PrismaClient", day: str) -> bool: + try: + await reconcile_day(prisma_client, day) + await _record_reconciled_through(prisma_client, day) + except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done + verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) + return False + return True + + +async def run_scheduled_daily_global_spend_reconcile( + prisma_client: "PrismaClient", + pod_lock_manager: "PodLockManager | None" = None, + alert: Callable[[str], Awaitable[None]] | None = None, + today: date | None = None, +) -> ReconcileResult | None: + """Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves + effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping.""" + redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache + if pod_lock_manager is None or redis_cache is None: + return await _run_and_alert(prisma_client, alert=alert, today=today) + + acquired: Final = await pod_lock_manager.acquire_lock( + cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS + ) + if not acquired and await _lock_is_held(pod_lock_manager, redis_cache): + verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run") + return None + try: + return await _run_and_alert(prisma_client, alert=alert, today=today) + finally: + if acquired: + await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) + + +async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool: + try: + lock_key: Final = pod_lock_manager.get_redis_lock_key(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) + return bool(await redis_cache.async_get_cache(lock_key)) + except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the run + verbose_proxy_logger.warning("Daily global spend reconcile: could not read the lock: %s", exc) + return False + + +async def _run_and_alert( + prisma_client: "PrismaClient", + *, + alert: Callable[[str], Awaitable[None]] | None, + today: date | None, +) -> ReconcileResult: + result: Final = await run_daily_global_spend_reconcile(prisma_client, today=today) + if result.days_reconciled: + verbose_proxy_logger.info( + "Daily global spend reconcile: rolled up %d day(s), reconciled through %s", + len(result.days_reconciled), + result.reconciled_through, + ) + if result.failed_day is not None and alert is not None: + await alert( + f"Daily global spend reconcile stopped at {result.failed_day}; usage totals keep reading the per-key " + f"table for ranges past {result.reconciled_through or 'the beginning'} until the next run succeeds." + ) + return result diff --git a/schema.prisma b/schema.prisma index 8072df5aa5b..5d433e916d6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -809,6 +809,35 @@ model LiteLLM_DailyUserSpend { @@index([endpoint]) } +// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view +model LiteLLM_DailyGlobalSpend { + id String @id @default(uuid()) + date String + model String? + model_group String? + custom_llm_provider String? + mcp_namespaced_tool_name String? + endpoint String? + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint]) + @@index([date]) +} + // Track daily organization spend metrics per model and key model LiteLLM_DailyOrganizationSpend { id String @id @default(uuid()) diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py index c1efb3e7220..cc443a2cfe5 100644 --- a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -1,12 +1,19 @@ """Tests for the single-statement daily spend upsert (LIT-5291).""" +import pathlib import re +from typing import Final +import psycopg import pytest +from psycopg.rows import dict_row +from pytest_postgresql import factories from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, + GLOBAL_SPEND_TABLE, build_bulk_upsert, + build_bulk_upsert_with_global_rollup, conflict_key, merge_by_conflict_key, ) @@ -185,3 +192,146 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null(): _, params = prisma_client.db.statements[0] assert None not in params[:9] assert transactions == {} + + +def user_txn(**overrides): + txn = {**tag_txn(), "user_id": "u-1", **overrides} + del txn["tag"] + del txn["request_id"] + return txn + + +def _bound_rows(insert_sql: str, params: tuple[object, ...]) -> list[dict[str, object]]: + """Each VALUES row of one INSERT as a column -> bound value mapping, consuming params in order.""" + header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", insert_sql) + assert header is not None, insert_sql + columns = [c.strip('"') for c in header.group(1).split(", ") if c != '"updated_at"'] + row_count = insert_sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") + return [dict(zip(columns, params[i * len(columns) : (i + 1) * len(columns)])) for i in range(row_count)] + + +def test_global_rollup_folds_every_key_and_user_into_one_row_per_dimension_tuple(): + """The global table has no api_key or user_id, so a batch spread over many keys and + users must collapse to one row per (date, model, group, provider, mcp, endpoint).""" + batch = merge_by_conflict_key( + USER_TABLE, + tuple(user_txn(user_id=f"u-{i}", api_key=f"sk-{i}", spend=1.0, api_requests=1) for i in range(5)) + + (user_txn(user_id="u-0", api_key="sk-0", model="claude", spend=10.0, api_requests=3),), + ) + + sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) + + entity_insert, global_insert = sql.split("RETURNING 1)") + entity_rows = _bound_rows(entity_insert, params) + global_rows = _bound_rows(global_insert, params[len(entity_rows) * len(entity_rows[0]) :]) + assert len(entity_rows) == 6 + assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in global_insert + assert [(r["model"], r["spend"], r["api_requests"]) for r in global_rows] == [ + ("claude", 10.0, 3), + ("gpt-4o-mini", 5.0, 5), + ] + assert all("api_key" not in r and "user_id" not in r for r in global_rows) + conflict = re.search(r"ON CONFLICT \(([^)]*)\)", global_insert) + assert conflict is not None + assert conflict.group(1) == ", ".join(f'"{c}"' for c in GLOBAL_SPEND_TABLE.key_columns) + + +def test_global_rollup_params_follow_the_entity_params_in_one_placeholder_sequence(): + """Both inserts bind from one flat tuple, so the global arm's placeholders must start + exactly where the entity arm's stop or every value lands one column off.""" + batch = merge_by_conflict_key(USER_TABLE, (user_txn(),)) + + sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) + + placeholders = [int(n) for n in re.findall(r"\$(\d+)::", sql)] + assert placeholders == list(range(1, len(params) + 1)) + + +_bulk_upsert_postgresql_proc: Final = factories.postgresql_proc() +_bulk_upsert_postgresql: Final = factories.postgresql("_bulk_upsert_postgresql_proc") + +_MIGRATIONS_DIR: Final = ( + pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" +) +_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" + +_DAILY_USER_SPEND_DDL: Final = """ + CREATE TABLE "LiteLLM_DailyUserSpend" ( + id TEXT PRIMARY KEY, + user_id TEXT, + date TEXT NOT NULL, + api_key TEXT NOT NULL, + model TEXT, + model_group TEXT, + custom_llm_provider TEXT, + mcp_namespaced_tool_name TEXT, + endpoint TEXT, + prompt_tokens BIGINT DEFAULT 0, + completion_tokens BIGINT DEFAULT 0, + cache_read_input_tokens BIGINT DEFAULT 0, + cache_creation_input_tokens BIGINT DEFAULT 0, + compression_saved_tokens BIGINT DEFAULT 0, + compression_savings_spend DOUBLE PRECISION DEFAULT 0, + prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, + spend DOUBLE PRECISION DEFAULT 0, + api_requests BIGINT DEFAULT 0, + successful_requests BIGINT DEFAULT 0, + failed_requests BIGINT DEFAULT 0, + created_at TIMESTAMP DEFAULT now(), + updated_at TIMESTAMP, + UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) + ) +""" + + +def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: + converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) + conn.execute( + converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + conn.commit() + + +def test_global_rollup_equals_the_per_key_sums_after_repeated_flushes(_bulk_upsert_postgresql: psycopg.Connection): + """Against real Postgres and the shipped migration: two flushes of a mixed batch leave + the global table exactly equal to the per-key table summed over user and key, with the + NULL and '' spellings of a dimension folded into one row.""" + conn: Final = _bulk_upsert_postgresql + conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal + conn.commit() + + batch = merge_by_conflict_key( + USER_TABLE, + ( + user_txn(user_id="u-1", api_key="sk-1", spend=1.0, prompt_tokens=10), + user_txn(user_id="u-2", api_key="sk-2", spend=2.0, prompt_tokens=20), + user_txn(user_id="u-1", api_key="sk-3", model=None, custom_llm_provider=None, spend=4.0), + user_txn(user_id="u-3", api_key="sk-4", model="", custom_llm_provider="", spend=8.0), + ), + ) + sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) + _execute_dollar_sql(conn, sql, params) + _execute_dollar_sql(conn, sql, params) + + with conn.cursor(row_factory=dict_row) as cur: + global_rows = cur.execute( + 'SELECT model, spend, prompt_tokens, api_requests FROM "LiteLLM_DailyGlobalSpend" ORDER BY model' + ).fetchall() + per_key = cur.execute( + """ + SELECT COALESCE(model, '') AS model, SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, + SUM(api_requests) AS api_requests + FROM "LiteLLM_DailyUserSpend" GROUP BY COALESCE(model, '') ORDER BY 1 + """ + ).fetchall() + + assert [row["model"] for row in global_rows] == ["", "gpt-4o-mini"] + assert [(r["model"], r["spend"], int(r["prompt_tokens"]), int(r["api_requests"])) for r in global_rows] == [ + (r["model"], float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"])) for r in per_key + ] + assert global_rows[0]["spend"] == pytest.approx(24.0) + assert global_rows[1]["spend"] == pytest.approx(6.0) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 5e977712a1e..d8a9013398e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -254,14 +254,19 @@ class _RecordingPrisma: def _row_values(statement: Statement, column: str) -> list[object]: - """Every row's value for one column, read out of the flat parameter tuple.""" + """Every row's value for one column of the first INSERT, read out of the flat parameter tuple. + + The user-table statement chains a global rollup INSERT after its own, so the row count + comes from the first INSERT's VALUES rather than from the parameter count. + """ sql, params = statement header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) assert header is not None, sql columns = header.group(1).split(", ") stride = len(columns) - 1 # updated_at is inlined, not bound offset = columns.index(f'"{column}"') - return [params[row * stride + offset] for row in range(len(params) // stride)] + rows = sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") + return [params[row * stride + offset] for row in range(rows)] @pytest.mark.asyncio @@ -1463,6 +1468,98 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected +def _entity_txn(entity_field: str, entity_id: str, api_key: str) -> dict[str, object]: + txn = _daily_txn() + del txn["user_id"] + return {**txn, entity_field: entity_id, "api_key": api_key} + + +@pytest.mark.asyncio +async def test_user_flush_writes_the_global_rollup_in_the_same_statement(): + """The user flush is the one place per-key spend becomes key-free spend, so a batch spread + over many keys must land in LiteLLM_DailyGlobalSpend as one row in the same statement. + A separate statement would let a crash between the two leave the tables out of sync.""" + prisma_client = _RecordingPrisma() + txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(4)} + + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions=txns, + entity_type="user", + entity_id_field="user_id", + ) + + assert len(prisma_client.db.statements) == 1 + sql, params = prisma_client.db.statements[0] + assert sql.count('INSERT INTO "LiteLLM_DailyUserSpend"') == 1 + assert sql.count('INSERT INTO "LiteLLM_DailyGlobalSpend"') == 1 + assert sql.index('"LiteLLM_DailyUserSpend"') < sql.index('"LiteLLM_DailyGlobalSpend"') + global_insert = sql.split('INSERT INTO "LiteLLM_DailyGlobalSpend"', 1)[1] + assert global_insert.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") == 1 + assert "api_key" not in global_insert + assert params.count(0.4) == 1 + assert txns == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entity_type", "entity_field"), + [ + ("team", "team_id"), + ("org", "organization_id"), + ("tag", "tag"), + ("end_user", "end_user_id"), + ("agent", "agent_id"), + ], +) +async def test_other_entity_flushes_leave_the_global_table_alone(entity_type, entity_field): + """Every entity table sees the same request, so writing the rollup from more than one of + them would count each request once per entity type.""" + prisma_client = _RecordingPrisma() + txn = _entity_txn(entity_field, "e-1", "sk-1") + if entity_type == "tag": + txn["request_id"] = "req-1" + + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions={"k": txn}, + entity_type=entity_type, + entity_id_field=entity_field, + ) + + (sql, _params) = prisma_client.db.statements[0] + assert "LiteLLM_DailyGlobalSpend" not in sql + + +@pytest.mark.asyncio +async def test_a_failed_chained_user_flush_keeps_every_transaction_for_retry(): + def raise_outage(): + raise ValueError("simulated database outage") + + prisma_client = _RecordingPrisma(execute_raw=raise_outage) + txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(3)} + expected = dict(txns) + mock_proxy_logging = MagicMock() + mock_proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(ValueError, match="simulated database outage"): + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=mock_proxy_logging, + daily_spend_transactions=txns, + entity_type="user", + entity_id_field="user_id", + ) + + assert txns == expected + assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in prisma_client.db.statements[0][0] + + @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 5ff3f89343b..5c74facae6a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -13,7 +13,13 @@ from pytest_postgresql import factories from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT +import pathlib + +from litellm.constants import ( + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, + PTU_SENTINEL_API_KEY, + USAGE_TOP_API_KEYS_LIMIT, +) from litellm.proxy.management_endpoints.common_daily_activity import ( _adjust_dates_for_timezone, _build_aggregated_sql_query, @@ -23,8 +29,11 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, + key_free_source_table, update_metrics, ) +from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL +from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, SpendMetrics, @@ -1618,6 +1627,149 @@ async def test_get_daily_activity_aggregated_explicit_api_key_filter_scopes_both assert set(day.breakdown.models["gpt-5"].api_key_breakdown) == {"key-1"} +def _prisma_with_marker(marker: str | None) -> MagicMock: + prisma = MagicMock() + prisma.db = MagicMock() + prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + row = None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + prisma.get_generic_data = AsyncMock(return_value=row) + return prisma + + +def _unfiltered_user_query(**overrides): + return { + "table_name": "litellm_dailyuserspend", + "entity_id_field": "user_id", + "entity_id": None, + "start_date": "2026-06-01", + "end_date": "2026-06-02", + "model": None, + "api_key": None, + "exclude_entity_ids": None, + "timezone_offset_minutes": None, + "include_current_utc_day": False, + **overrides, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("marker", "overrides", "expected"), + [ + ("2026-06-02", {}, "LiteLLM_DailyGlobalSpend"), + ("2026-06-02", {"model": "gpt-5"}, "LiteLLM_DailyGlobalSpend"), + ("2026-06-01", {}, None), + (None, {}, None), + ("2026-06-02", {"api_key": "sk-1"}, None), + ("2026-06-02", {"api_key": []}, None), + ("2026-06-02", {"entity_id": "u-1"}, None), + ("2026-06-02", {"exclude_entity_ids": ["u-1"]}, None), + ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), + ], +) +async def test_key_free_source_table_routes_only_unfiltered_user_reads_within_the_marker(marker, overrides, expected): + """Anything that filters by key or entity has no counterpart in the global table, and a + range the reconcile has not reached must stay on the per-key table.""" + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + prisma = _prisma_with_marker(marker) + + assert await key_free_source_table(prisma, _unfiltered_user_query(**overrides)) == expected + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +@pytest.mark.asyncio +async def test_key_free_source_table_judges_the_timezone_extended_end_not_the_requested_one(): + """A caller west of UTC asking through their local today gets today's UTC bucket added to + the range; the marker must cover that extended day, not just the requested end.""" + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + today_utc: Final = datetime.now(timezone.utc).date() + yesterday: Final = (today_utc - timedelta(days=1)).isoformat() + query: Final = _unfiltered_user_query( + start_date=yesterday, end_date=yesterday, timezone_offset_minutes=24 * 60, include_current_utc_day=True + ) + + assert await key_free_source_table(_prisma_with_marker(yesterday), query) is None + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + assert await key_free_source_table(_prisma_with_marker(today_utc.isoformat()), query) == "LiteLLM_DailyGlobalSpend" + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +_GLOBAL_SPEND_MIGRATION: Final = ( + pathlib.Path(__file__).resolve().parents[4] + / "litellm-proxy-extras" + / "litellm_proxy_extras" + / "migrations" + / "20260915000000_add_daily_global_spend" + / "migration.sql" +) + + +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_free_arm( + _aggregated_postgresql: psycopg.Connection, +): + """With the range reconciled, the key-free arm reads LiteLLM_DailyGlobalSpend while the + per-key arm stays on the user table, and the response is identical to the all-per-key + read: same totals, same rollups, same top keys.""" + n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 + rows: Final = [ + ( + f"row-{day}-{i:03d}", + f"user-{i % 7}", + day, + f"key-{i:03d}", + "gpt-5" if i % 2 else "claude", + "" if i % 3 else "gpt-5", + "openai" if i % 2 else None, + "/v1/chat/completions" if i % 5 else None, + 10, + float(i + 1), + 1, + 1, + ) + for day in ("2026-06-01", "2026-06-02") + for i in range(n_keys) + ] + _seed_daily_user_spend(_aggregated_postgresql, rows) + with _aggregated_postgresql.cursor() as cur: + cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal + for day in ("2026-06-01", "2026-06-02"): + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg + {"p1": day}, + ) + _aggregated_postgresql.commit() + + async def read(marker: str | None, sql_seen: list[str]): + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + prisma = _prisma_with_marker(marker) + run_query = _psycopg_query_raw(_aggregated_postgresql, []) + + async def query_raw(sql: str, *params: str): + sql_seen.append(sql) + return await run_query(sql, *params) + + prisma.db.query_raw = query_raw + return await get_daily_activity_aggregated( + prisma_client=prisma, + entity_metadata_field=None, + **_unfiltered_user_query(), + ) + + per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim + global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim + from_per_key = await read(None, per_key_sql) + from_global = await read("2026-06-02", global_sql) + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 + assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 + assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 2 + assert from_global.model_dump() == from_per_key.model_dump() + assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT + assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace( diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index deb7289d2d1..ee72e98ffa9 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -1042,6 +1042,54 @@ async def test_spend_report_locks_are_never_released(): proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() +def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, MagicMock]: + scheduler = MagicMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.alerting_handler = AsyncMock() + prisma_client = MagicMock() + ProxyStartupEvent._initialize_daily_global_spend_reconcile_job( + scheduler=scheduler, + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma_client, + ) + return scheduler, proxy_logging_obj, prisma_client + + +def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run(): + """Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a + fresh deploy switches usage reads to the global table without waiting for the nightly + run, and replaces any previous registration of the same job id.""" + from datetime import datetime, timedelta, timezone + + from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID + + scheduler, _, _ = _init_daily_global_spend_reconcile_job() + + (call,) = scheduler.add_job.call_args_list + assert call.kwargs["id"] == DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID + assert call.kwargs["replace_existing"] is True + assert call.args[1:] == ("cron",) + assert (call.kwargs["hour"], call.kwargs["minute"], call.kwargs["timezone"]) == (0, 30, "UTC") + assert timedelta(0) < call.kwargs["next_run_time"] - datetime.now(timezone.utc) <= timedelta(minutes=2) + + +@pytest.mark.asyncio +async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch): + scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job() + run = AsyncMock() + monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run) + + await scheduler.add_job.call_args.args[0]() + + run.assert_awaited_once() + assert run.await_args.args == (prisma_client,) + assert run.await_args.kwargs["pod_lock_manager"] is proxy_logging_obj.db_spend_update_writer.pod_lock_manager + await run.await_args.kwargs["alert"]("day 2026-09-01 failed") + proxy_logging_obj.alerting_handler.assert_awaited_once() + assert proxy_logging_obj.alerting_handler.await_args.kwargs["message"] == "day 2026-09-01 failed" + assert proxy_logging_obj.alerting_handler.await_args.kwargs["level"] == "High" + + @pytest.mark.asyncio async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch): """The boot-time send goes through the same gate, so a losing pod sends nothing at all: diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py new file mode 100644 index 00000000000..13dc757cbbd --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -0,0 +1,382 @@ +"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" + +import pathlib +import re +from contextlib import asynccontextmanager +from datetime import date +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import psycopg +import pytest +from psycopg.rows import dict_row +from pytest_postgresql import factories + +from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + build_bulk_upsert_with_global_rollup, + merge_by_conflict_key, +) +from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + RECONCILE_DAY_SQL, + reconciled_through, + run_daily_global_spend_reconcile, + run_scheduled_daily_global_spend_reconcile, +) +from litellm.proxy.utils import evict_config_param + +USER_TABLE: Final = DAILY_SPEND_TABLES["user"] +TODAY: Final = date(2026, 9, 15) + + +class _FakeConfigRow: + def __init__(self, param_name: str, param_value: object) -> None: + self.param_name = param_name + self.param_value = param_value + + +class _FakeConfigTable: + def __init__(self) -> None: + self.rows: dict[str, object] = {} + + async def upsert(self, *, where: dict[str, str], data: dict[str, dict[str, str]]) -> _FakeConfigRow: + self.rows[where["param_name"]] = data["update"]["param_value"] + return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + + +class _FakeTransaction: + def __init__(self, prisma: "_FakePrisma") -> None: + self._prisma = prisma + + async def execute_raw(self, sql: str, *params: str) -> int: + if "LOCK TABLE" in sql: + self._prisma.locks_taken += 1 + return 0 + (day,) = params + if day in self._prisma.failing_days: + raise RuntimeError(f"day {day} exploded") + self._prisma.reconciled.append(day) + return 1 + + +class _FakeDb: + def __init__(self, prisma: "_FakePrisma") -> None: + self._prisma = prisma + self.litellm_config = _FakeConfigTable() + + async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: + first, last = params + return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] + + @asynccontextmanager + async def tx(self, timeout: object): + yield _FakeTransaction(self._prisma) + + +class _FakePrisma: + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and a transaction.""" + + def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: + self.user_days = user_days + self.failing_days = failing_days + self.reconciled: list[str] = [] + self.locks_taken = 0 + self.db = _FakeDb(self) + + async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: + stored = self.db.litellm_config.rows.get(value) + return None if stored is None else _FakeConfigRow(value, stored) + + +@pytest.fixture(autouse=True) +async def _fresh_marker_cache(): + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + yield + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + +@pytest.mark.asyncio +async def test_first_run_rolls_up_every_historical_day_and_today_then_marks_today(): + """Before any marker exists, every day with per-key rows is rolled up, plus today even + with no rows yet, so reads for ranges ending today can switch to the global table.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14")) + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15") + assert result.failed_day is None + assert result.reconciled_through == "2026-09-15" + assert await reconciled_through(prisma) == "2026-09-15" + assert prisma.locks_taken == 4 + + +@pytest.mark.asyncio +async def test_later_run_replays_the_marker_day_and_the_day_before_only(): + """Days older than marker-1 are settled; the marker day and its predecessor are replayed so + rows a pre-writer pod flushed around midnight during a rolling deploy get folded in.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14", "2026-09-15") + assert "2026-09-01" not in prisma.reconciled + assert await reconciled_through(prisma) == "2026-09-15" + + +@pytest.mark.asyncio +async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_good_day(): + """The marker may never claim a day that was not rewritten: reads past it would then trust + a global table missing that day's spend.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01",) + assert result.failed_day == "2026-09-02" + assert result.reconciled_through == "2026-09-01" + assert prisma.reconciled == ["2026-09-01"] + assert await reconciled_through(prisma) == "2026-09-01" + + +@pytest.mark.asyncio +async def test_the_next_run_resumes_from_the_failed_day(): + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) + await run_daily_global_spend_reconcile(prisma, today=TODAY) + prisma.failing_days = frozenset() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03", "2026-09-15") + assert await reconciled_through(prisma) == "2026-09-15" + + +@pytest.mark.asyncio +async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): + """A pre-writer pod flushing rows for the day before the marker is exactly the replay case; + when that replay fails the marker must stay put and the operator must hear about it.""" + prisma = _FakePrisma(user_days=("2026-09-13",)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + prisma.user_days = ("2026-09-12", "2026-09-13") + prisma.failing_days = frozenset({"2026-09-12"}) + alert = AsyncMock() + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + + assert result is not None + assert result.days_reconciled == () + assert result.failed_day == "2026-09-12" + assert result.reconciled_through == "2026-09-13" + alert.assert_awaited_once() + assert "2026-09-12" in alert.await_args.args[0] + + +@pytest.mark.asyncio +async def test_a_clean_run_does_not_alert(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + alert = AsyncMock() + + await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + + alert.assert_not_awaited() + + +def _pod_lock(acquired: bool) -> MagicMock: + lock = MagicMock() + lock.redis_cache = MagicMock() + lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod") + lock.get_redis_lock_key = MagicMock(return_value="lock-key") + lock.acquire_lock = AsyncMock(return_value=acquired) + lock.release_lock = AsyncMock() + return lock + + +@pytest.mark.asyncio +async def test_scheduled_run_skips_when_another_pod_holds_the_lock(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + lock = _pod_lock(acquired=False) + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + + assert result is None + assert prisma.reconciled == [] + lock.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + lock = _pod_lock(acquired=True) + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + + assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + lock.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read(): + """A Redis outage must not stall the backfill: the day rewrite is idempotent, so running + twice is only wasted effort while skipping forever leaves usage on the slow path.""" + prisma = _FakePrisma(user_days=("2026-09-13",)) + lock = _pod_lock(acquired=False) + lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + + assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + lock.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_marker_is_read_back_from_the_json_string_the_config_table_stores(): + prisma = _FakePrisma(user_days=()) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-10"}' + + assert await reconciled_through(prisma) == "2026-09-10" + + +@pytest.mark.asyncio +async def test_an_unparseable_marker_reads_as_never_reconciled(): + prisma = _FakePrisma(user_days=()) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"something_else": 1}' + + assert await reconciled_through(prisma) is None + + +_rollup_postgresql_proc: Final = factories.postgresql_proc() +_rollup_postgresql: Final = factories.postgresql("_rollup_postgresql_proc") + +_MIGRATIONS_DIR: Final = ( + pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" +) +_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" + +_DAILY_USER_SPEND_DDL: Final = """ + CREATE TABLE "LiteLLM_DailyUserSpend" ( + id TEXT PRIMARY KEY, + user_id TEXT, + date TEXT NOT NULL, + api_key TEXT NOT NULL, + model TEXT, + model_group TEXT, + custom_llm_provider TEXT, + mcp_namespaced_tool_name TEXT, + endpoint TEXT, + prompt_tokens BIGINT DEFAULT 0, + completion_tokens BIGINT DEFAULT 0, + cache_read_input_tokens BIGINT DEFAULT 0, + cache_creation_input_tokens BIGINT DEFAULT 0, + compression_saved_tokens BIGINT DEFAULT 0, + compression_savings_spend DOUBLE PRECISION DEFAULT 0, + prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, + spend DOUBLE PRECISION DEFAULT 0, + api_requests BIGINT DEFAULT 0, + successful_requests BIGINT DEFAULT 0, + failed_requests BIGINT DEFAULT 0, + created_at TIMESTAMP DEFAULT now(), + updated_at TIMESTAMP, + UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) + ) +""" + +_PER_KEY_SUMS_SQL: Final = """ + SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group, + COALESCE(custom_llm_provider, '') AS custom_llm_provider, + SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests + FROM "LiteLLM_DailyUserSpend" WHERE date = %s + GROUP BY 1, 2, 3 ORDER BY 1, 2, 3 +""" +_GLOBAL_ROWS_SQL: Final = """ + SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests + FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3 +""" + + +def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: + converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) + conn.execute( + converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query + {f"p{i}": v for i, v in enumerate(params, start=1)}, + ) + conn.commit() + + +def _user_txn(**overrides): + return { + "user_id": "u-1", + "date": "2026-09-14", + "api_key": "sk-1", + "model": "gpt-5", + "model_group": "gpt-5", + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": "", + "endpoint": "/chat/completions", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 1.0, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + **overrides, + } + + +def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: + return [ + ( + r["model"], + r["model_group"], + r["custom_llm_provider"], + float(r["spend"]), + int(r["prompt_tokens"]), + int(r["api_requests"]), + ) # pyright: ignore[reportArgumentType] # dict_row values are untyped + for r in rows + ] + + +def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): + """Against real Postgres and the shipped migration: rows the writer never saw (a + pre-writer pod's flush, NULL and '' dimension spellings) end up folded into the global + day, running the day twice changes nothing, and other days are left alone.""" + conn: Final = _rollup_postgresql + conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal + conn.commit() + + written_batch = merge_by_conflict_key( + USER_TABLE, + (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), + ) + _execute_dollar_sql(conn, *build_bulk_upsert_with_global_rollup(USER_TABLE, written_batch)) + + conn.execute( + """ + INSERT INTO "LiteLLM_DailyUserSpend" + (id, user_id, date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, + endpoint, prompt_tokens, spend, api_requests) + VALUES + ('legacy-1', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', NULL, 'openai', NULL, NULL, 5, 4.0, 1), + ('legacy-2', 'u-9', '2026-09-14', 'sk-9', 'gpt-5', '', 'openai', '', '', 5, 8.0, 1), + ('legacy-3', 'u-9', '2026-09-13', 'sk-9', 'claude', '', 'anthropic', '', '', 7, 16.0, 1) + """ + ) + conn.commit() + + _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) + _execute_dollar_sql(conn, RECONCILE_DAY_SQL, ("2026-09-14",)) + + with conn.cursor(row_factory=dict_row) as cur: + global_rows = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-14",)).fetchall() + per_key = cur.execute(_PER_KEY_SUMS_SQL, ("2026-09-14",)).fetchall() + untouched = cur.execute(_GLOBAL_ROWS_SQL, ("2026-09-13",)).fetchall() + + assert _normalized(global_rows) == _normalized(per_key) + assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped + assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] + assert untouched == [] From ad8de0e1927c18d5d14c92939bbd531c54573874 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 23:42:19 +0000 Subject: [PATCH 2/8] fix(proxy): roll up only closed days into LiteLLM_DailyGlobalSpend and split the key-free read at the marker The write path no longer dual-writes the global table. The cron rolls up closed UTC days only, so a pod still flushing the current day can never leave the global table short. The key-free arm reads days through the marker from the global table and later days from LiteLLM_DailyUserSpend in one UNION ALL, and the marker comes from the config cache rather than a per-request database lookup. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/daily_spend_bulk_upsert.py | 98 +++--------- litellm/proxy/db/db_spend_update_writer.py | 7 +- .../common_daily_activity.py | 79 ++++++--- .../daily_global_spend_rollup.py | 43 ++--- .../proxy/db/test_daily_spend_bulk_upsert.py | 150 ------------------ .../proxy/db/test_db_spend_update_writer.py | 101 +----------- .../test_common_daily_activity.py | 90 ++++++----- .../test_daily_global_spend_rollup.py | 92 +++++------ 8 files changed, 204 insertions(+), 456 deletions(-) diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py index c83043101eb..a143643577e 100644 --- a/litellm/proxy/db/daily_spend_bulk_upsert.py +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -25,41 +25,29 @@ SpendRow = Mapping[str, object] @dataclass(frozen=True, slots=True) class DailySpendTable: - """A daily rollup table and the unique constraint its upserts arbitrate on.""" + """The physical table behind one entity's daily rollup.""" name: str - key_columns: tuple[str, ...] + entity_id_column: str carries_request_id: bool = False +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), + "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), + "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), + "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), + "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), + "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + } +) + # The unique constraint's columns after the entity id, in constraint order. A NULL can # never match itself in a unique index, so every one of these is normalized to '': the # conflict target has to be NULL-free or the row is re-inserted on every single flush. _KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") - -def _entity_table(name: str, entity_id_column: str, carries_request_id: bool = False) -> DailySpendTable: - return DailySpendTable( - name=name, key_columns=(entity_id_column, *_KEY_COLUMNS), carries_request_id=carries_request_id - ) - - -DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( - { - "user": _entity_table("LiteLLM_DailyUserSpend", "user_id"), - "team": _entity_table("LiteLLM_DailyTeamSpend", "team_id"), - "org": _entity_table("LiteLLM_DailyOrganizationSpend", "organization_id"), - "end_user": _entity_table("LiteLLM_DailyEndUserSpend", "end_user_id"), - "agent": _entity_table("LiteLLM_DailyAgentSpend", "agent_id"), - "tag": _entity_table("LiteLLM_DailyTagSpend", "tag", carries_request_id=True), - } -) - -GLOBAL_SPEND_TABLE: Final = DailySpendTable( - name="LiteLLM_DailyGlobalSpend", - key_columns=("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"), -) - _COUNTER_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -104,7 +92,7 @@ def _as_float(value: object) -> float: def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" - return tuple(_as_text(transaction.get(column)) for column in table.key_columns) + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) def _merge(group: Sequence[SpendRow]) -> SpendRow: @@ -142,11 +130,7 @@ def _row_params( return ( str(uuid.uuid4()), *key, - *( - () - if "model_group" in table.key_columns - else (None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")),) - ), + None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), @@ -156,25 +140,26 @@ def _row_params( def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: return ( "id", - *table.key_columns, - *(() if "model_group" in table.key_columns else ("model_group",)), + table.entity_id_column, + *_KEY_COLUMNS, + "model_group", *_COUNTER_COLUMNS, *_SPEND_COLUMNS, *(("request_id",) if table.carries_request_id else ()), ) -def _upsert_statement( +def build_bulk_upsert( table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]], - first_param: int, -) -> str: +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" columns: Final = _insert_columns(table) quoted_table: Final = f'"{table.name}"' rows: Final = ", ".join( "(" + ", ".join( - f"${first_param + row_index * len(columns) + offset}::{_CASTS.get(column, 'text')}" + f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" for offset, column in enumerate(columns) ) + ", (NOW() AT TIME ZONE 'UTC'))" @@ -191,44 +176,11 @@ def _upsert_statement( if table.carries_request_id else "" ) - return ( + sql: Final = ( f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' f"VALUES {rows}\n" - f"ON CONFLICT ({_quoted(table.key_columns)}) DO UPDATE SET\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" f" {increments}{request_id_update},\n" f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) - - -def _params(table: DailySpendTable, batch: Sequence[tuple[tuple[str, ...], SpendRow]]) -> tuple[SqlValue, ...]: - return tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) - - -def build_bulk_upsert( - table: DailySpendTable, - batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """The single statement writing one merged batch, plus its positional arguments.""" - return _upsert_statement(table, batch, first_param=1), _params(table, batch) - - -def build_bulk_upsert_with_global_rollup( - table: DailySpendTable, - batch: Sequence[tuple[tuple[str, ...], SpendRow]], -) -> tuple[str, tuple[SqlValue, ...]]: - """One statement writing a batch to its table and, atomically, its key-free rollup - to ``LiteLLM_DailyGlobalSpend``. - - A data-modifying CTE runs both inserts in the same snapshot and transaction, so a - batch that lands in one table lands in both and a retried deadlock replays both. - Postgres does not order the CTE against the main statement, so two writers can still - deadlock across the tables; the caller's deadlock retry covers that, and each insert - takes its own rows in key order so same-table lock order stays deterministic. - """ - global_batch: Final = merge_by_conflict_key(GLOBAL_SPEND_TABLE, tuple(row for _, row in batch)) - entity_params: Final = _params(table, batch) - sql: Final = ( - f"WITH entity_rows AS (\n{_upsert_statement(table, batch, first_param=1)}\nRETURNING 1)\n" - f"{_upsert_statement(GLOBAL_SPEND_TABLE, global_batch, first_param=len(entity_params) + 1)}" - ) - return sql, (*entity_params, *_params(GLOBAL_SPEND_TABLE, global_batch)) + return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d5c839a9be8..eaa03c5d7f7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -46,7 +46,6 @@ from litellm.proxy._types import ( from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, build_bulk_upsert, - build_bulk_upsert_with_global_rollup, merge_by_conflict_key, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -1940,11 +1939,7 @@ class DBSpendUpdateWriter: merged_batch = merge_by_conflict_key( table=table, transactions=tuple(transactions_to_process.values()) ) - sql, params = ( - build_bulk_upsert_with_global_rollup(table=table, batch=merged_batch) - if entity_type == "user" - else build_bulk_upsert(table=table, batch=merged_batch) - ) + sql, params = build_bulk_upsert(table=table, batch=merged_batch) await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f1d78dca201..b90f874c04c 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -11,8 +11,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT from litellm.proxy._types import CommonProxyErrors -from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE -from litellm.proxy.spend_tracking.daily_global_spend_rollup import reconciled_through +from litellm.proxy.spend_tracking.daily_global_spend_rollup import GLOBAL_SPEND_TABLE_NAME, reconciled_through from litellm.proxy.spend_tracking.key_metadata_recovery import ( attach_user_emails, recover_double_hashed_key_metadata, @@ -736,28 +735,62 @@ def _rollup_metric_select(table_name: str) -> str: _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" -async def key_free_source_table(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: - """The table the key-free arm reads from, when the global rollup can answer instead of the per-key table. +_KEY_FREE_SOURCE_COLUMNS: Final = ( + "date", + "model", + "model_group", + "custom_llm_provider", + "mcp_namespaced_tool_name", + "endpoint", + "spend", + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", + "compression_savings_spend", + "prompt_caching_savings_spend", + "gateway_injected_caching_savings_spend", + "autorouter_savings_spend", + "api_requests", + "successful_requests", + "failed_requests", +) - Only an unfiltered read of the user table has the same rows as ``LiteLLM_DailyGlobalSpend``, - and only through the day the reconcile marker has reached: the writer keeps that day - current, later days are covered once the next run advances the marker. + +async def global_rollup_reconciled_through(prisma_client: PrismaClient, query: _AggregatedQueryKwargs) -> str | None: + """The last day ``LiteLLM_DailyGlobalSpend`` can answer the key-free arm for, or None to + read it all from the per-key table. + + Only an unfiltered read of the user table sums to the same rows as the global table. The + marker read is served from the config cache, so this is not a database round trip per request. """ if query["table_name"] != "litellm_dailyuserspend": return None if query["entity_id"] is not None or query["api_key"] is not None or query["exclude_entity_ids"]: return None - _, adjusted_end = _adjust_dates_for_timezone( - query["start_date"], query["end_date"], query["timezone_offset_minutes"], query["include_current_utc_day"] - ) try: - marker: Final = await reconciled_through(prisma_client) + return await reconciled_through(prisma_client) except Exception as exc: # noqa: BLE001 # the per-key table is always a correct answer, so never fail the read verbose_proxy_logger.warning("Could not read the daily global spend marker, using the per-key table: %s", exc) return None - if marker is None or adjusted_end > marker: - return None - return GLOBAL_SPEND_TABLE.name + + +def _key_free_source(pg_table: str, where_clause: str, marker_param: str | None) -> str: + """The relation the key-free arm aggregates: the per-key table alone, or the global rollup + for days through the marker plus the per-key table for the days still open after it.""" + if marker_param is None: + return f'"{pg_table}"\n WHERE {where_clause}' + columns: Final = ", ".join(_KEY_FREE_SOURCE_COLUMNS) + return f"""( + SELECT {columns} + FROM "{GLOBAL_SPEND_TABLE_NAME}" + WHERE {where_clause} AND date <= {marker_param} + UNION ALL + SELECT {columns} + FROM "{pg_table}" + WHERE {where_clause} AND date > {marker_param} + ) AS key_free_source""" def _build_aggregated_sql_query( @@ -772,15 +805,16 @@ def _build_aggregated_sql_query( exclude_entity_ids: list[str] | None = None, # mutable-ok: filter union shared with the paginated path timezone_offset_minutes: int | None = None, include_current_utc_day: bool = False, - key_free_table: str | None = None, + global_rollup_through: str | None = None, ) -> tuple[str, list[str]]: # mutable-ok: SQL text plus its ordered $N params """Build the GROUPING SETS query for aggregated daily activity. One statement, two UNION ALL arms over the same WHERE clause. The first arm is key-free: grand total, per-date totals and the (date, model / model_group / provider / mcp / endpoint) rollups, so its row count never grows with the number - of keys; it reads ``key_free_table`` when given (the global rollup, whose row count - never grew with the number of keys to begin with) and the entity table otherwise. + of keys. With ``global_rollup_through`` it reads days through that marker from + ``LiteLLM_DailyGlobalSpend`` (whose row count never grew with the number of keys to + begin with) and only the days after it from the per-key table. The second arm emits the (date, , api_key) rollups for the USAGE_TOP_API_KEYS_LIMIT highest-spend keys only. Both arms share the 7-bit group_level bitmask (date, api_key, model, model_group, provider, mcp, endpoint). @@ -806,8 +840,8 @@ def _build_aggregated_sql_query( exclude_entity_ids=exclude_entity_ids, ) sentinel_param: Final = f"${len(where_params) + 1}" + marker_param: Final = None if global_rollup_through is None else f"${len(where_params) + 2}" metric_select: Final = _rollup_metric_select(table_name) - key_free_source: Final = key_free_table or pg_table # TODO: drop the successful_requests/failed_requests aggregates (and the # total_successful_requests metadata they feed) once the admin UI reads SGR @@ -826,8 +860,7 @@ def _build_aggregated_sql_query( | GROUPING(model, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level,{metric_select} - FROM "{key_free_source}" - WHERE {where_clause} + FROM {_key_free_source(pg_table, where_clause, marker_param)} GROUP BY GROUPING SETS ( (date), (date, model), @@ -869,7 +902,8 @@ def _build_aggregated_sql_query( )) """ - return sql_query, [*where_params, PTU_SENTINEL_API_KEY] + marker_params: Final = () if global_rollup_through is None else (global_rollup_through,) + return sql_query, [*where_params, PTU_SENTINEL_API_KEY, *marker_params] def _build_entity_rollup_sql_query( @@ -1418,7 +1452,8 @@ async def get_daily_activity_aggregated( include_current_utc_day=include_current_utc_day, ) sql_query, sql_params = _build_aggregated_sql_query( - **query_kwargs, key_free_table=await key_free_source_table(prisma_client, query_kwargs) + **query_kwargs, + global_rollup_through=await global_rollup_reconciled_through(prisma_client, query_kwargs), ) entity_query: Final = _build_entity_rollup_sql_query(**query_kwargs) if include_entity_breakdown else None diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 9d344421332..a9fb7669785 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -1,9 +1,10 @@ -"""Reconcile ``LiteLLM_DailyGlobalSpend`` from ``LiteLLM_DailyUserSpend``, one day per transaction. +"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``. -The spend writer keeps both tables in step from the moment it is deployed; this job rolls up -the days before that and records how far it has reached in ``LiteLLM_Config`` so usage reads -know when the global table can answer for a date range. It runs as a background cron, never -in a Prisma migration, since on a large deployment the aggregate is minutes of work. +Only days that are over get rolled up, so a pod still flushing per-key spend for the current +day can never leave the global table short; usage reads serve days through the recorded +marker from the global table and later days live from the per-key table. The marker lives in +``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on a +large deployment the first backfill is minutes of work. """ from collections.abc import Awaitable, Callable @@ -19,7 +20,6 @@ from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ) -from litellm.proxy.db.daily_spend_bulk_upsert import GLOBAL_SPEND_TABLE from litellm.repositories.config_repository import ConfigRepository if TYPE_CHECKING: @@ -27,8 +27,11 @@ if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient -_DAY_TRANSACTION_TIMEOUT: Final = timedelta(minutes=10) _REPLAY_DAYS: Final = 1 +GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" +# The unique constraint, in constraint order. NULL never matches itself in a unique index, so +# every column is normalized to '' or the same group would be inserted again on every run. +_KEY_COLUMNS: Final = ("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") _METRIC_COLUMNS: Final = ( "prompt_tokens", "completion_tokens", @@ -51,23 +54,21 @@ def _quoted(columns: tuple[str, ...]) -> str: def _reconcile_day_sql() -> str: - key_columns: Final = GLOBAL_SPEND_TABLE.key_columns - normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in key_columns) + normalized_keys: Final = ", ".join(f"COALESCE(\"{column}\", '')" for column in _KEY_COLUMNS) sums: Final = ", ".join(f'SUM("{column}")' for column in _METRIC_COLUMNS) overwrite: Final = ", ".join(f'"{column}" = EXCLUDED."{column}"' for column in _METRIC_COLUMNS) return ( - f'INSERT INTO "{GLOBAL_SPEND_TABLE.name}" ("id", {_quoted(key_columns)}, {_quoted(_METRIC_COLUMNS)}, ' + f'INSERT INTO "{GLOBAL_SPEND_TABLE_NAME}" ("id", {_quoted(_KEY_COLUMNS)}, {_quoted(_METRIC_COLUMNS)}, ' '"updated_at")\n' f"SELECT gen_random_uuid()::text, {normalized_keys}, {sums}, (NOW() AT TIME ZONE 'UTC')\n" 'FROM "LiteLLM_DailyUserSpend" WHERE "date" = $1\n' f"GROUP BY {normalized_keys}\n" - f"ON CONFLICT ({_quoted(key_columns)}) DO UPDATE SET {overwrite}, " + f"ON CONFLICT ({_quoted(_KEY_COLUMNS)}) DO UPDATE SET {overwrite}, " "\"updated_at\" = (NOW() AT TIME ZONE 'UTC')" ) RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_LOCK_GLOBAL_TABLE_SQL: Final = f'LOCK TABLE "{GLOBAL_SPEND_TABLE.name}" IN EXCLUSIVE MODE' _PENDING_DAYS_SQL: Final = ( 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' ) @@ -134,19 +135,19 @@ def _first_pending_day(marker: str | None) -> str: async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - """Every UTC day through today still to roll up, oldest first; the marker day and the one - before it are replayed so rows flushed by a pre-writer pod during a rolling deploy are folded in.""" + """Every closed UTC day (strictly before today) still to roll up, oldest first. The marker + day and the one before it are replayed so per-key rows that landed after their day was + rolled up (a flush straddling midnight, a late retry) are folded in.""" marker: Final = await reconciled_through(prisma_client) - rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), today.isoformat()) - return tuple(sorted({*(_DateRow.model_validate(row).date for row in rows), today.isoformat()})) + last_closed_day: Final = (today - timedelta(days=1)).isoformat() + rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), last_closed_day) + return tuple(_DateRow.model_validate(row).date for row in rows) async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: - """Rewrite one day of the global table from the per-key sums; the table lock keeps the - writer's increments out between the aggregate and the overwrite so none are lost.""" - async with prisma_client.db.tx(timeout=_DAY_TRANSACTION_TIMEOUT) as transaction: - await transaction.execute_raw(_LOCK_GLOBAL_TABLE_SQL) - await transaction.execute_raw(RECONCILE_DAY_SQL, day) + """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun + overwrites every group with the same totals.""" + await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) async def run_daily_global_spend_reconcile( diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py index cc443a2cfe5..c1efb3e7220 100644 --- a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -1,19 +1,12 @@ """Tests for the single-statement daily spend upsert (LIT-5291).""" -import pathlib import re -from typing import Final -import psycopg import pytest -from psycopg.rows import dict_row -from pytest_postgresql import factories from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, - GLOBAL_SPEND_TABLE, build_bulk_upsert, - build_bulk_upsert_with_global_rollup, conflict_key, merge_by_conflict_key, ) @@ -192,146 +185,3 @@ async def test_writer_survives_a_transaction_whose_key_columns_are_null(): _, params = prisma_client.db.statements[0] assert None not in params[:9] assert transactions == {} - - -def user_txn(**overrides): - txn = {**tag_txn(), "user_id": "u-1", **overrides} - del txn["tag"] - del txn["request_id"] - return txn - - -def _bound_rows(insert_sql: str, params: tuple[object, ...]) -> list[dict[str, object]]: - """Each VALUES row of one INSERT as a column -> bound value mapping, consuming params in order.""" - header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", insert_sql) - assert header is not None, insert_sql - columns = [c.strip('"') for c in header.group(1).split(", ") if c != '"updated_at"'] - row_count = insert_sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") - return [dict(zip(columns, params[i * len(columns) : (i + 1) * len(columns)])) for i in range(row_count)] - - -def test_global_rollup_folds_every_key_and_user_into_one_row_per_dimension_tuple(): - """The global table has no api_key or user_id, so a batch spread over many keys and - users must collapse to one row per (date, model, group, provider, mcp, endpoint).""" - batch = merge_by_conflict_key( - USER_TABLE, - tuple(user_txn(user_id=f"u-{i}", api_key=f"sk-{i}", spend=1.0, api_requests=1) for i in range(5)) - + (user_txn(user_id="u-0", api_key="sk-0", model="claude", spend=10.0, api_requests=3),), - ) - - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - - entity_insert, global_insert = sql.split("RETURNING 1)") - entity_rows = _bound_rows(entity_insert, params) - global_rows = _bound_rows(global_insert, params[len(entity_rows) * len(entity_rows[0]) :]) - assert len(entity_rows) == 6 - assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in global_insert - assert [(r["model"], r["spend"], r["api_requests"]) for r in global_rows] == [ - ("claude", 10.0, 3), - ("gpt-4o-mini", 5.0, 5), - ] - assert all("api_key" not in r and "user_id" not in r for r in global_rows) - conflict = re.search(r"ON CONFLICT \(([^)]*)\)", global_insert) - assert conflict is not None - assert conflict.group(1) == ", ".join(f'"{c}"' for c in GLOBAL_SPEND_TABLE.key_columns) - - -def test_global_rollup_params_follow_the_entity_params_in_one_placeholder_sequence(): - """Both inserts bind from one flat tuple, so the global arm's placeholders must start - exactly where the entity arm's stop or every value lands one column off.""" - batch = merge_by_conflict_key(USER_TABLE, (user_txn(),)) - - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - - placeholders = [int(n) for n in re.findall(r"\$(\d+)::", sql)] - assert placeholders == list(range(1, len(params) + 1)) - - -_bulk_upsert_postgresql_proc: Final = factories.postgresql_proc() -_bulk_upsert_postgresql: Final = factories.postgresql("_bulk_upsert_postgresql_proc") - -_MIGRATIONS_DIR: Final = ( - pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" -) -_GLOBAL_SPEND_MIGRATION: Final = _MIGRATIONS_DIR / "20260915000000_add_daily_global_spend" / "migration.sql" - -_DAILY_USER_SPEND_DDL: Final = """ - CREATE TABLE "LiteLLM_DailyUserSpend" ( - id TEXT PRIMARY KEY, - user_id TEXT, - date TEXT NOT NULL, - api_key TEXT NOT NULL, - model TEXT, - model_group TEXT, - custom_llm_provider TEXT, - mcp_namespaced_tool_name TEXT, - endpoint TEXT, - prompt_tokens BIGINT DEFAULT 0, - completion_tokens BIGINT DEFAULT 0, - cache_read_input_tokens BIGINT DEFAULT 0, - cache_creation_input_tokens BIGINT DEFAULT 0, - compression_saved_tokens BIGINT DEFAULT 0, - compression_savings_spend DOUBLE PRECISION DEFAULT 0, - prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, - autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, - spend DOUBLE PRECISION DEFAULT 0, - api_requests BIGINT DEFAULT 0, - successful_requests BIGINT DEFAULT 0, - failed_requests BIGINT DEFAULT 0, - created_at TIMESTAMP DEFAULT now(), - updated_at TIMESTAMP, - UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) - ) -""" - - -def _execute_dollar_sql(conn: psycopg.Connection, sql: str, params: tuple[object, ...]) -> None: - converted: Final = re.sub(r"\$(\d+)", r"%(p\1)s", sql) - conn.execute( - converted, # pyright: ignore[reportArgumentType] # psycopg stubs want a literal-typed query - {f"p{i}": v for i, v in enumerate(params, start=1)}, - ) - conn.commit() - - -def test_global_rollup_equals_the_per_key_sums_after_repeated_flushes(_bulk_upsert_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: two flushes of a mixed batch leave - the global table exactly equal to the per-key table summed over user and key, with the - NULL and '' spellings of a dimension folded into one row.""" - conn: Final = _bulk_upsert_postgresql - conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal - conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - conn.commit() - - batch = merge_by_conflict_key( - USER_TABLE, - ( - user_txn(user_id="u-1", api_key="sk-1", spend=1.0, prompt_tokens=10), - user_txn(user_id="u-2", api_key="sk-2", spend=2.0, prompt_tokens=20), - user_txn(user_id="u-1", api_key="sk-3", model=None, custom_llm_provider=None, spend=4.0), - user_txn(user_id="u-3", api_key="sk-4", model="", custom_llm_provider="", spend=8.0), - ), - ) - sql, params = build_bulk_upsert_with_global_rollup(USER_TABLE, batch) - _execute_dollar_sql(conn, sql, params) - _execute_dollar_sql(conn, sql, params) - - with conn.cursor(row_factory=dict_row) as cur: - global_rows = cur.execute( - 'SELECT model, spend, prompt_tokens, api_requests FROM "LiteLLM_DailyGlobalSpend" ORDER BY model' - ).fetchall() - per_key = cur.execute( - """ - SELECT COALESCE(model, '') AS model, SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, - SUM(api_requests) AS api_requests - FROM "LiteLLM_DailyUserSpend" GROUP BY COALESCE(model, '') ORDER BY 1 - """ - ).fetchall() - - assert [row["model"] for row in global_rows] == ["", "gpt-4o-mini"] - assert [(r["model"], r["spend"], int(r["prompt_tokens"]), int(r["api_requests"])) for r in global_rows] == [ - (r["model"], float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"])) for r in per_key - ] - assert global_rows[0]["spend"] == pytest.approx(24.0) - assert global_rows[1]["spend"] == pytest.approx(6.0) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index d8a9013398e..5e977712a1e 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -254,19 +254,14 @@ class _RecordingPrisma: def _row_values(statement: Statement, column: str) -> list[object]: - """Every row's value for one column of the first INSERT, read out of the flat parameter tuple. - - The user-table statement chains a global rollup INSERT after its own, so the row count - comes from the first INSERT's VALUES rather than from the parameter count. - """ + """Every row's value for one column, read out of the flat parameter tuple.""" sql, params = statement header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) assert header is not None, sql columns = header.group(1).split(", ") stride = len(columns) - 1 # updated_at is inlined, not bound offset = columns.index(f'"{column}"') - rows = sql.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") - return [params[row * stride + offset] for row in range(rows)] + return [params[row * stride + offset] for row in range(len(params) // stride)] @pytest.mark.asyncio @@ -1468,98 +1463,6 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry(): assert daily_spend_transactions == expected -def _entity_txn(entity_field: str, entity_id: str, api_key: str) -> dict[str, object]: - txn = _daily_txn() - del txn["user_id"] - return {**txn, entity_field: entity_id, "api_key": api_key} - - -@pytest.mark.asyncio -async def test_user_flush_writes_the_global_rollup_in_the_same_statement(): - """The user flush is the one place per-key spend becomes key-free spend, so a batch spread - over many keys must land in LiteLLM_DailyGlobalSpend as one row in the same statement. - A separate statement would let a crash between the two leave the tables out of sync.""" - prisma_client = _RecordingPrisma() - txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(4)} - - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=MagicMock(), - daily_spend_transactions=txns, - entity_type="user", - entity_id_field="user_id", - ) - - assert len(prisma_client.db.statements) == 1 - sql, params = prisma_client.db.statements[0] - assert sql.count('INSERT INTO "LiteLLM_DailyUserSpend"') == 1 - assert sql.count('INSERT INTO "LiteLLM_DailyGlobalSpend"') == 1 - assert sql.index('"LiteLLM_DailyUserSpend"') < sql.index('"LiteLLM_DailyGlobalSpend"') - global_insert = sql.split('INSERT INTO "LiteLLM_DailyGlobalSpend"', 1)[1] - assert global_insert.split("ON CONFLICT", 1)[0].count("(NOW() AT TIME ZONE 'UTC'))") == 1 - assert "api_key" not in global_insert - assert params.count(0.4) == 1 - assert txns == {} - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("entity_type", "entity_field"), - [ - ("team", "team_id"), - ("org", "organization_id"), - ("tag", "tag"), - ("end_user", "end_user_id"), - ("agent", "agent_id"), - ], -) -async def test_other_entity_flushes_leave_the_global_table_alone(entity_type, entity_field): - """Every entity table sees the same request, so writing the rollup from more than one of - them would count each request once per entity type.""" - prisma_client = _RecordingPrisma() - txn = _entity_txn(entity_field, "e-1", "sk-1") - if entity_type == "tag": - txn["request_id"] = "req-1" - - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=MagicMock(), - daily_spend_transactions={"k": txn}, - entity_type=entity_type, - entity_id_field=entity_field, - ) - - (sql, _params) = prisma_client.db.statements[0] - assert "LiteLLM_DailyGlobalSpend" not in sql - - -@pytest.mark.asyncio -async def test_a_failed_chained_user_flush_keeps_every_transaction_for_retry(): - def raise_outage(): - raise ValueError("simulated database outage") - - prisma_client = _RecordingPrisma(execute_raw=raise_outage) - txns = {f"k{i}": _entity_txn("user_id", f"user-{i}", f"sk-{i}") for i in range(3)} - expected = dict(txns) - mock_proxy_logging = MagicMock() - mock_proxy_logging.failure_handler = AsyncMock() - - with pytest.raises(ValueError, match="simulated database outage"): - await DBSpendUpdateWriter._update_daily_spend( - n_retry_times=0, - prisma_client=prisma_client, - proxy_logging_obj=mock_proxy_logging, - daily_spend_transactions=txns, - entity_type="user", - entity_id_field="user_id", - ) - - assert txns == expected - assert 'INSERT INTO "LiteLLM_DailyGlobalSpend"' in prisma_client.db.statements[0][0] - - @pytest.mark.asyncio async def test_commit_key_spend_updates_includes_last_active(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 5c74facae6a..12e5fe6af4d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1,3 +1,4 @@ +import pathlib import re from collections.abc import Sequence from datetime import datetime, timedelta, timezone @@ -10,11 +11,6 @@ import pytest from psycopg.rows import dict_row from pytest_postgresql import factories -from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR - - -import pathlib - from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, PTU_SENTINEL_API_KEY, @@ -29,10 +25,11 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, - key_free_source_table, + global_rollup_reconciled_through, update_metrics, ) from litellm.proxy.spend_tracking.daily_global_spend_rollup import RECONCILE_DAY_SQL +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import evict_config_param from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, @@ -1632,7 +1629,9 @@ def _prisma_with_marker(marker: str | None) -> MagicMock: prisma.db = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) - row = None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + row = ( + None if marker is None else SimpleNamespace(param_name="m", param_value=f'{{"reconciled_through": "{marker}"}}') + ) prisma.get_generic_data = AsyncMock(return_value=row) return prisma @@ -1657,9 +1656,9 @@ def _unfiltered_user_query(**overrides): @pytest.mark.parametrize( ("marker", "overrides", "expected"), [ - ("2026-06-02", {}, "LiteLLM_DailyGlobalSpend"), - ("2026-06-02", {"model": "gpt-5"}, "LiteLLM_DailyGlobalSpend"), - ("2026-06-01", {}, None), + ("2026-06-02", {}, "2026-06-02"), + ("2026-06-02", {"model": "gpt-5"}, "2026-06-02"), + ("2026-05-01", {}, "2026-05-01"), (None, {}, None), ("2026-06-02", {"api_key": "sk-1"}, None), ("2026-06-02", {"api_key": []}, None), @@ -1668,33 +1667,50 @@ def _unfiltered_user_query(**overrides): ("2026-06-02", {"table_name": "litellm_dailyteamspend", "entity_id_field": "team_id"}, None), ], ) -async def test_key_free_source_table_routes_only_unfiltered_user_reads_within_the_marker(marker, overrides, expected): - """Anything that filters by key or entity has no counterpart in the global table, and a - range the reconcile has not reached must stay on the per-key table.""" +async def test_global_rollup_marker_is_used_only_for_unfiltered_user_reads(marker, overrides, expected): + """Anything that filters by key or entity has no counterpart in the global table; the + SQL splits the range at the marker itself, so the marker passes through unchanged.""" await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) prisma = _prisma_with_marker(marker) - assert await key_free_source_table(prisma, _unfiltered_user_query(**overrides)) == expected + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query(**overrides)) == expected await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) @pytest.mark.asyncio -async def test_key_free_source_table_judges_the_timezone_extended_end_not_the_requested_one(): - """A caller west of UTC asking through their local today gets today's UTC bucket added to - the range; the marker must cover that extended day, not just the requested end.""" +async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table(): await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - today_utc: Final = datetime.now(timezone.utc).date() - yesterday: Final = (today_utc - timedelta(days=1)).isoformat() - query: Final = _unfiltered_user_query( - start_date=yesterday, end_date=yesterday, timezone_offset_minutes=24 * 60, include_current_utc_day=True - ) + prisma = _prisma_with_marker(None) + prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down")) - assert await key_free_source_table(_prisma_with_marker(yesterday), query) is None - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - assert await key_free_source_table(_prisma_with_marker(today_utc.isoformat()), query) == "LiteLLM_DailyGlobalSpend" + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) +def test_aggregated_sql_splits_the_key_free_arm_at_the_marker_and_keeps_the_key_arm_per_key(): + sql, params = _build_aggregated_sql_query(**_unfiltered_user_query(), global_rollup_through="2026-06-01") + marker_param: Final = f"${len(params)}" + + assert params[-1] == "2026-06-01" + assert ( + f'FROM "LiteLLM_DailyGlobalSpend"\n WHERE date >= $1 AND date <= $2 AND date <= {marker_param}' + in sql + ) + assert ( + f'FROM "LiteLLM_DailyUserSpend"\n WHERE date >= $1 AND date <= $2 AND date > {marker_param}' in sql + ) + key_arm: Final = sql.split("UNION ALL\n (WITH top_api_keys")[1] + assert "LiteLLM_DailyGlobalSpend" not in key_arm + assert marker_param not in key_arm + + +def test_aggregated_sql_without_a_marker_reads_the_per_key_table_only(): + sql, params = _build_aggregated_sql_query(**_unfiltered_user_query()) + + assert "LiteLLM_DailyGlobalSpend" not in sql + assert params[-1] == PTU_SENTINEL_API_KEY + + _GLOBAL_SPEND_MIGRATION: Final = ( pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" @@ -1706,12 +1722,12 @@ _GLOBAL_SPEND_MIGRATION: Final = ( @pytest.mark.asyncio -async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_free_arm( +async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_table_and_open_days_live( _aggregated_postgresql: psycopg.Connection, ): - """With the range reconciled, the key-free arm reads LiteLLM_DailyGlobalSpend while the - per-key arm stays on the user table, and the response is identical to the all-per-key - read: same totals, same rollups, same top keys.""" + """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must + give the same response as reading everything per-key: day 1 from the global table, day 2 + live, one grand total across both. The per-key arm stays on the user table throughout.""" n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 rows: Final = [ ( @@ -1734,11 +1750,10 @@ async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_ _seed_daily_user_spend(_aggregated_postgresql, rows) with _aggregated_postgresql.cursor() as cur: cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal - for day in ("2026-06-01", "2026-06-02"): - cur.execute( - re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg - {"p1": day}, - ) + cur.execute( + re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg + {"p1": "2026-06-01"}, + ) _aggregated_postgresql.commit() async def read(marker: str | None, sql_seen: list[str]): @@ -1760,16 +1775,19 @@ async def test_get_daily_activity_aggregated_reads_the_global_table_for_the_key_ per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim from_per_key = await read(None, per_key_sql) - from_global = await read("2026-06-02", global_sql) + from_global = await read("2026-06-01", global_sql) await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 - assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 2 + assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} + + def _no_spend_record(): """A rollup row for a key with no spend, where SUM() returns NULL (None).""" return SimpleNamespace( diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 13dc757cbbd..11ca72e7b3d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -2,7 +2,6 @@ import pathlib import re -from contextlib import asynccontextmanager from datetime import date from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -13,11 +12,7 @@ from psycopg.rows import dict_row from pytest_postgresql import factories from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM -from litellm.proxy.db.daily_spend_bulk_upsert import ( - DAILY_SPEND_TABLES, - build_bulk_upsert_with_global_rollup, - merge_by_conflict_key, -) +from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( RECONCILE_DAY_SQL, reconciled_through, @@ -45,21 +40,6 @@ class _FakeConfigTable: return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) -class _FakeTransaction: - def __init__(self, prisma: "_FakePrisma") -> None: - self._prisma = prisma - - async def execute_raw(self, sql: str, *params: str) -> int: - if "LOCK TABLE" in sql: - self._prisma.locks_taken += 1 - return 0 - (day,) = params - if day in self._prisma.failing_days: - raise RuntimeError(f"day {day} exploded") - self._prisma.reconciled.append(day) - return 1 - - class _FakeDb: def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma @@ -69,19 +49,21 @@ class _FakeDb: first, last = params return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] - @asynccontextmanager - async def tx(self, timeout: object): - yield _FakeTransaction(self._prisma) + async def execute_raw(self, sql: str, *params: str) -> int: + (day,) = params + if day in self._prisma.failing_days: + raise RuntimeError(f"day {day} exploded") + self._prisma.reconciled.append(day) + return 1 class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and a transaction.""" + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: self.user_days = user_days self.failing_days = failing_days self.reconciled: list[str] = [] - self.locks_taken = 0 self.db = _FakeDb(self) async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: @@ -97,33 +79,45 @@ async def _fresh_marker_cache(): @pytest.mark.asyncio -async def test_first_run_rolls_up_every_historical_day_and_today_then_marks_today(): - """Before any marker exists, every day with per-key rows is rolled up, plus today even - with no rows yet, so reads for ranges ending today can switch to the global table.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14")) +async def test_first_run_rolls_up_every_closed_day_and_never_today(): + """Before any marker exists every closed day with per-key rows is rolled up. Today is left + out: pods are still flushing it, so it is served live from the per-key table until it closes.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15") + assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") assert result.failed_day is None - assert result.reconciled_through == "2026-09-15" - assert await reconciled_through(prisma) == "2026-09-15" - assert prisma.locks_taken == 4 + assert result.reconciled_through == "2026-09-14" + assert await reconciled_through(prisma) == "2026-09-14" + assert "2026-09-15" not in prisma.reconciled @pytest.mark.asyncio async def test_later_run_replays_the_marker_day_and_the_day_before_only(): """Days older than marker-1 are settled; the marker day and its predecessor are replayed so - rows a pre-writer pod flushed around midnight during a rolling deploy get folded in.""" + per-key rows that landed after their day was rolled up get folded in.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.reconciled.clear() result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14", "2026-09-15") + assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14") assert "2026-09-01" not in prisma.reconciled - assert await reconciled_through(prisma) == "2026-09-15" + assert await reconciled_through(prisma) == "2026-09-14" + + +@pytest.mark.asyncio +async def test_a_run_with_no_new_closed_days_keeps_the_marker(): + prisma = _FakePrisma(user_days=("2026-09-13",)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + + assert result.days_reconciled == ("2026-09-13",) + assert result.reconciled_through == "2026-09-13" @pytest.mark.asyncio @@ -149,16 +143,16 @@ async def test_the_next_run_resumes_from_the_failed_day(): result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03", "2026-09-15") - assert await reconciled_through(prisma) == "2026-09-15" + assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") + assert await reconciled_through(prisma) == "2026-09-03" @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """A pre-writer pod flushing rows for the day before the marker is exactly the replay case; - when that replay fails the marker must stay put and the operator must hear about it.""" + """A late flush for the day before the marker is exactly the replay case; when that replay + fails the marker must stay put and the operator must hear about it.""" prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 13)) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.user_days = ("2026-09-12", "2026-09-13") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() @@ -212,7 +206,7 @@ async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) - assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_awaited_once() @@ -226,7 +220,7 @@ async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read() result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) - assert result is not None and result.days_reconciled == ("2026-09-13", "2026-09-15") + assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_not_awaited() @@ -341,9 +335,9 @@ def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_postgresql: psycopg.Connection): - """Against real Postgres and the shipped migration: rows the writer never saw (a - pre-writer pod's flush, NULL and '' dimension spellings) end up folded into the global - day, running the day twice changes nothing, and other days are left alone.""" + """Against real Postgres and the shipped migration: writer-shaped rows and legacy rows + (NULL and '' dimension spellings) fold into one global day, running the day twice changes + nothing, and other days are left alone.""" conn: Final = _rollup_postgresql conn.execute(_DAILY_USER_SPEND_DDL) # pyright: ignore[reportArgumentType] # DDL literal conn.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal @@ -353,7 +347,7 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p USER_TABLE, (_user_txn(api_key="sk-1", spend=1.0), _user_txn(api_key="sk-2", user_id="u-2", spend=2.0, prompt_tokens=20)), ) - _execute_dollar_sql(conn, *build_bulk_upsert_with_global_rollup(USER_TABLE, written_batch)) + _execute_dollar_sql(conn, *build_bulk_upsert(USER_TABLE, written_batch)) conn.execute( """ From 84c098df92f8d89ed5d083466ec62347110062e0 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 00:42:03 +0000 Subject: [PATCH 3/8] fix(proxy): fold late-arriving per-key spend into already rolled-up global days The reconcile now records the database clock of the scan behind the last complete run and, on the next run, rewrites every closed day with per-key rows updated since then, however old the day is. Replaying only the marker day and the one before it missed a delayed flush or retry that landed on an older date, and reads through the marker come from the global table alone, so that spend was never counted. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 128 +++++++++++++----- .../test_daily_global_spend_rollup.py | 88 ++++++++++-- 2 files changed, 168 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index a9fb7669785..73068381dab 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -2,9 +2,12 @@ Only days that are over get rolled up, so a pod still flushing per-key spend for the current day can never leave the global table short; usage reads serve days through the recorded -marker from the global table and later days live from the per-key table. The marker lives in -``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on a -large deployment the first backfill is minutes of work. +marker from the global table and later days live from the per-key table. Per-key rows are +dated by request start, so spend can land on a day that was already rolled up (a flush +straddling midnight, a retry after an outage). Each run therefore also rewrites every closed +day that has rows touched since the previous run's scan, whatever the date. The marker lives +in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on +a large deployment the first backfill is minutes of work. """ from collections.abc import Awaitable, Callable @@ -27,7 +30,6 @@ if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient -_REPLAY_DAYS: Final = 1 GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" # The unique constraint, in constraint order. NULL never matches itself in a unique index, so # every column is normalized to '' or the same group would be inserted again on every run. @@ -69,15 +71,26 @@ def _reconcile_day_sql() -> str: RECONCILE_DAY_SQL: Final = _reconcile_day_sql() +_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now" +_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' +# Pod clocks drift from the database clock and from each other, so rows are picked up from a +# little before the previous scan; rewriting a day twice is idempotent. _PENDING_DAYS_SQL: Final = ( - 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' + 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ' + 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' + 'ORDER BY "date"' ) class ReconciledThrough(BaseModel): + """``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is + the database clock when the scan behind the last fully successful run started: every per-key + row written before it, on any day through the marker, is in the global table.""" + model_config = ConfigDict(frozen=True, extra="ignore") reconciled_through: str + scanned_at: str | None = None class _MarkerRow(BaseModel): @@ -92,6 +105,12 @@ class _DateRow(BaseModel): date: str +class _NowRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + now: str + + @dataclass(frozen=True, slots=True) class ReconcileResult: days_reconciled: tuple[str, ...] @@ -99,49 +118,70 @@ class ReconcileResult: failed_day: str | None = None -def _marker_from_param_value(value: object) -> str | None: +@dataclass(frozen=True, slots=True) +class _PendingScan: + marker: ReconciledThrough | None + scanned_at: str + days: tuple[str, ...] + + +def _marker_from_param_value(value: object) -> ReconciledThrough | None: try: - parsed: Final = ( + return ( ReconciledThrough.model_validate_json(value) if isinstance(value, str) else ReconciledThrough.model_validate(value) ) except ValidationError: return None - return parsed.reconciled_through -async def reconciled_through(prisma_client: "PrismaClient") -> str | None: - """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" +async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: from litellm.proxy.utils import get_config_param row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) -async def _record_reconciled_through(prisma_client: "PrismaClient", day: str) -> None: +async def reconciled_through(prisma_client: "PrismaClient") -> str | None: + """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" + marker: Final = await read_marker(prisma_client) + return None if marker is None else marker.reconciled_through + + +async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThrough) -> None: from litellm.proxy.utils import invalidate_config_param await ConfigRepository(prisma_client).set_param( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ReconciledThrough(reconciled_through=day).model_dump_json() + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, marker.model_dump_json() ) await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -def _first_pending_day(marker: str | None) -> str: - if marker is None: - return "" - return (date.fromisoformat(marker) - timedelta(days=_REPLAY_DAYS)).isoformat() +async def _db_now(prisma_client: "PrismaClient") -> str: + rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) + return _NowRow.model_validate(rows[0]).now + + +async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingScan: + """Every closed UTC day (strictly before today) still to roll up, oldest first: days past the + marker, plus any day with per-key rows written since the scan behind the marker. Before a + run has fully succeeded there is no such scan, so every closed day is rolled up.""" + marker: Final = await read_marker(prisma_client) + scanned_at: Final = await _db_now(prisma_client) + last_closed_day: Final = (today - timedelta(days=1)).isoformat() + rows: Final = ( + await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) + if marker is None or marker.scanned_at is None + else await prisma_client.db.query_raw( + _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at + ) + ) + return _PendingScan(marker, scanned_at, tuple(_DateRow.model_validate(row).date for row in rows)) async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - """Every closed UTC day (strictly before today) still to roll up, oldest first. The marker - day and the one before it are replayed so per-key rows that landed after their day was - rolled up (a flush straddling midnight, a late retry) are folded in.""" - marker: Final = await reconciled_through(prisma_client) - last_closed_day: Final = (today - timedelta(days=1)).isoformat() - rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), last_closed_day) - return tuple(_DateRow.model_validate(row).date for row in rows) + return (await _scan_pending(prisma_client, today)).days async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: @@ -155,26 +195,42 @@ async def run_daily_global_spend_reconcile( today: date | None = None, ) -> ReconcileResult: """Roll up every pending day, advancing the marker after each; a failing day stops the run - with the marker on the last good day so the next run resumes there.""" + with the marker on the last good day so the next run resumes there. The scan time is only + recorded once every pending day is done, so late rows a failed run saw are found again.""" effective_today: Final = today or datetime.now(timezone.utc).date() - days: Final = await pending_days(prisma_client, effective_today) - done: Final = await _reconcile_until_failure(prisma_client, days) - failed: Final = days[len(done)] if len(done) < len(days) else None - marker: Final = done[-1] if done else await reconciled_through(prisma_client) - return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=failed) + scan: Final = await _scan_pending(prisma_client, effective_today) + done: Final = await _reconcile_until_failure(prisma_client, scan) + if len(done) < len(scan.days): + marker: Final = await reconciled_through(prisma_client) + return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) + if scan.marker is not None or done: + await _record_marker(prisma_client, _advanced(scan.marker, done, scanned_at=scan.scanned_at)) + return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) -async def _reconcile_until_failure(prisma_client: "PrismaClient", days: tuple[str, ...]) -> tuple[str, ...]: - for index, day in enumerate(days): - if not await _reconcile_and_record(prisma_client, day): - return days[:index] - return days +def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanned_at: str | None) -> ReconciledThrough: + """The marker after ``days`` were rewritten: a late old day never moves it back.""" + through: Final = max((marker.reconciled_through if marker is not None else "", *days)) + return ReconciledThrough(reconciled_through=through, scanned_at=scanned_at) -async def _reconcile_and_record(prisma_client: "PrismaClient", day: str) -> bool: +async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: + for index, day in enumerate(scan.days): + if not await _reconcile_and_record(prisma_client, scan.marker, scan.days[: index + 1]): + return scan.days[:index] + return scan.days + + +async def _reconcile_and_record( + prisma_client: "PrismaClient", marker: ReconciledThrough | None, done_with_this: tuple[str, ...] +) -> bool: + day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_reconciled_through(prisma_client, day) + await _record_marker( + prisma_client, + _advanced(marker, done_with_this, scanned_at=None if marker is None else marker.scanned_at), + ) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 11ca72e7b3d..9a098744f08 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -15,6 +15,7 @@ from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( RECONCILE_DAY_SQL, + read_marker, reconciled_through, run_daily_global_spend_reconcile, run_scheduled_daily_global_spend_reconcile, @@ -41,13 +42,25 @@ class _FakeConfigTable: class _FakeDb: + """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, + so "rows written since the last scan" behaves like Postgres would.""" + def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma self.litellm_config = _FakeConfigTable() async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: - first, last = params - return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] + if sql.startswith("SELECT (NOW()"): + self._prisma.clock += 1 + return [{"now": f"clock-{self._prisma.clock:04d}"}] + rows = self._prisma.user_rows + if len(params) == 1: + (last,) = params + return [{"date": d} for d in sorted(rows) if d <= last] + last, marker, scanned_at = params + return [ + {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) + ] async def execute_raw(self, sql: str, *params: str) -> int: (day,) = params @@ -61,11 +74,17 @@ class _FakePrisma: """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: - self.user_days = user_days + self.clock = 0 + self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days self.reconciled: list[str] = [] self.db = _FakeDb(self) + def write_late_row(self, day: str) -> None: + """A per-key row for ``day`` lands now, after whatever scans already happened.""" + self.clock += 1 + self.user_rows[day] = f"clock-{self.clock:04d}" + async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: stored = self.db.litellm_config.rows.get(value) return None if stored is None else _FakeConfigRow(value, stored) @@ -94,20 +113,66 @@ async def test_first_run_rolls_up_every_closed_day_and_never_today(): @pytest.mark.asyncio -async def test_later_run_replays_the_marker_day_and_the_day_before_only(): - """Days older than marker-1 are settled; the marker day and its predecessor are replayed so - per-key rows that landed after their day was rolled up get folded in.""" +async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.reconciled.clear() result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14") - assert "2026-09-01" not in prisma.reconciled + assert result.days_reconciled == ("2026-09-14",) assert await reconciled_through(prisma) == "2026-09-14" +@pytest.mark.asyncio +async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): + """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a + day far behind the marker. That day is rewritten, and the marker never moves back for it.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.reconciled.clear() + prisma.write_late_row("2026-09-01") + prisma.write_late_row("2026-09-03") + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-03") + assert "2026-09-05" not in prisma.reconciled + assert await reconciled_through(prisma) == "2026-09-13" + + +@pytest.mark.asyncio +async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): + """The scan time only advances when every pending day was rewritten, otherwise a late row + found by the failed run would be counted as handled.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.write_late_row("2026-09-01") + prisma.failing_days = frozenset({"2026-09-01"}) + failed = await run_daily_global_spend_reconcile(prisma, today=TODAY) + prisma.failing_days = frozenset() + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert failed.failed_day == "2026-09-01" + assert failed.reconciled_through == "2026-09-13" + assert result.days_reconciled == ("2026-09-01",) + assert result.failed_day is None + + +@pytest.mark.asyncio +async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-13") + marker = await read_marker(prisma) + assert marker is not None and marker.reconciled_through == "2026-09-13" and marker.scanned_at is not None + + @pytest.mark.asyncio async def test_a_run_with_no_new_closed_days_keeps_the_marker(): prisma = _FakePrisma(user_days=("2026-09-13",)) @@ -116,7 +181,7 @@ async def test_a_run_with_no_new_closed_days_keeps_the_marker(): result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) - assert result.days_reconciled == ("2026-09-13",) + assert result.days_reconciled == () assert result.reconciled_through == "2026-09-13" @@ -149,11 +214,10 @@ async def test_the_next_run_resumes_from_the_failed_day(): @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """A late flush for the day before the marker is exactly the replay case; when that replay - fails the marker must stay put and the operator must hear about it.""" + """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" prisma = _FakePrisma(user_days=("2026-09-13",)) await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) - prisma.user_days = ("2026-09-12", "2026-09-13") + prisma.write_late_row("2026-09-12") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() From 0601d2bb03646c596a680d580b0f9bb5a83ee237 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 01:58:37 +0000 Subject: [PATCH 4/8] feat(proxy): carry response time metrics through LiteLLM_DailyGlobalSpend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 2 ++ .../litellm_proxy_extras/schema.prisma | 2 ++ .../management_endpoints/common_daily_activity.py | 2 ++ litellm/proxy/schema.prisma | 2 ++ .../spend_tracking/daily_global_spend_rollup.py | 2 ++ schema.prisma | 2 ++ .../test_common_daily_activity.py | 6 ++++++ .../test_daily_global_spend_rollup.py | 13 +++++++++++-- 8 files changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql index 1d6cdea0c7b..d0bc3e159de 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql @@ -20,6 +20,8 @@ CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGlobalSpend" ( "api_requests" BIGINT NOT NULL DEFAULT 0, "successful_requests" BIGINT NOT NULL DEFAULT 0, "failed_requests" BIGINT NOT NULL DEFAULT 0, + "total_response_time_ms" BIGINT NOT NULL DEFAULT 0, + "timed_requests" BIGINT NOT NULL DEFAULT 0, "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, "updated_at" TIMESTAMP(3) NOT NULL, diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a73e8774c87..42769c323a9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -837,6 +837,8 @@ model LiteLLM_DailyGlobalSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 31ec0c51b7d..47465324f42 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -767,6 +767,8 @@ _KEY_FREE_SOURCE_COLUMNS: Final = ( "api_requests", "successful_requests", "failed_requests", + "total_response_time_ms", + "timed_requests", ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a73e8774c87..42769c323a9 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -837,6 +837,8 @@ model LiteLLM_DailyGlobalSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 73068381dab..ba4a4e4e3d6 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -43,6 +43,8 @@ _METRIC_COLUMNS: Final = ( "api_requests", "successful_requests", "failed_requests", + "total_response_time_ms", + "timed_requests", "compression_savings_spend", "prompt_caching_savings_spend", "gateway_injected_caching_savings_spend", diff --git a/schema.prisma b/schema.prisma index a73e8774c87..42769c323a9 100644 --- a/schema.prisma +++ b/schema.prisma @@ -837,6 +837,8 @@ model LiteLLM_DailyGlobalSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 2f25c507e0f..ee9f886acec 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1771,6 +1771,10 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ ] _seed_daily_user_spend(_aggregated_postgresql, rows) with _aggregated_postgresql.cursor() as cur: + cur.execute( + 'UPDATE "LiteLLM_DailyUserSpend" SET total_response_time_ms = prompt_tokens * 25, ' + "timed_requests = api_requests" + ) cur.execute(_GLOBAL_SPEND_MIGRATION.read_text()) # pyright: ignore[reportArgumentType] # DDL literal cur.execute( re.sub(r"\$(\d+)", r"%(p\1)s", RECONCILE_DAY_SQL), # pyright: ignore[reportArgumentType] # $N -> psycopg @@ -1805,6 +1809,8 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25 + assert from_global.metadata.total_timed_requests == 2 * n_keys assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 9a098744f08..b5b78229c11 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -336,6 +336,8 @@ _DAILY_USER_SPEND_DDL: Final = """ api_requests BIGINT DEFAULT 0, successful_requests BIGINT DEFAULT 0, failed_requests BIGINT DEFAULT 0, + total_response_time_ms BIGINT DEFAULT 0, + timed_requests BIGINT DEFAULT 0, created_at TIMESTAMP DEFAULT now(), updated_at TIMESTAMP, UNIQUE (user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint) @@ -345,12 +347,14 @@ _DAILY_USER_SPEND_DDL: Final = """ _PER_KEY_SUMS_SQL: Final = """ SELECT COALESCE(model, '') AS model, COALESCE(model_group, '') AS model_group, COALESCE(custom_llm_provider, '') AS custom_llm_provider, - SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests + SUM(spend) AS spend, SUM(prompt_tokens) AS prompt_tokens, SUM(api_requests) AS api_requests, + SUM(total_response_time_ms) AS total_response_time_ms, SUM(timed_requests) AS timed_requests FROM "LiteLLM_DailyUserSpend" WHERE date = %s GROUP BY 1, 2, 3 ORDER BY 1, 2, 3 """ _GLOBAL_ROWS_SQL: Final = """ - SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests + SELECT model, model_group, custom_llm_provider, spend, prompt_tokens, api_requests, + total_response_time_ms, timed_requests FROM "LiteLLM_DailyGlobalSpend" WHERE date = %s ORDER BY 1, 2, 3 """ @@ -380,6 +384,8 @@ def _user_txn(**overrides): "api_requests": 1, "successful_requests": 1, "failed_requests": 0, + "total_response_time_ms": 800, + "timed_requests": 1, **overrides, } @@ -393,6 +399,8 @@ def _normalized(rows: list[dict[str, object]]) -> list[tuple[object, ...]]: float(r["spend"]), int(r["prompt_tokens"]), int(r["api_requests"]), + int(r["total_response_time_ms"]), + int(r["timed_requests"]), ) # pyright: ignore[reportArgumentType] # dict_row values are untyped for r in rows ] @@ -436,5 +444,6 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p assert _normalized(global_rows) == _normalized(per_key) assert sum(float(r["spend"]) for r in global_rows) == pytest.approx(15.0) # pyright: ignore[reportArgumentType] # dict_row values are untyped + assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] assert untouched == [] From 834313af4b188ee5561e8c0e8094fad399e5ac78 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 12:13:57 +0000 Subject: [PATCH 5/8] test(proxy): assert the global rollup split and scheduler through behavior, not SQL text or add_job arguments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_common_daily_activity.py | 95 ++++++++++--------- .../proxy/proxy_server/test_lifecycle.py | 24 +++-- 2 files changed, 66 insertions(+), 53 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index ac761315a27..fc3ede88aa9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -1647,30 +1647,6 @@ async def test_global_rollup_marker_read_failure_falls_back_to_the_per_key_table await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -def test_aggregated_sql_splits_the_key_free_arm_at_the_marker_and_keeps_the_key_arm_per_key(): - sql, params = _build_aggregated_sql_query(**_unfiltered_user_query(), global_rollup_through="2026-06-01") - marker_param: Final = f"${len(params)}" - - assert params[-1] == "2026-06-01" - assert ( - f'FROM "LiteLLM_DailyGlobalSpend"\n WHERE date >= $1 AND date <= $2 AND date <= {marker_param}' - in sql - ) - assert ( - f'FROM "LiteLLM_DailyUserSpend"\n WHERE date >= $1 AND date <= $2 AND date > {marker_param}' in sql - ) - key_arm: Final = sql.split("UNION ALL\n (WITH top_api_keys")[1] - assert "LiteLLM_DailyGlobalSpend" not in key_arm - assert marker_param not in key_arm - - -def test_aggregated_sql_without_a_marker_reads_the_per_key_table_only(): - sql, params = _build_aggregated_sql_query(**_unfiltered_user_query()) - - assert "LiteLLM_DailyGlobalSpend" not in sql - assert params[-1] == PTU_SENTINEL_API_KEY - - _GLOBAL_SPEND_MIGRATION: Final = ( pathlib.Path(__file__).resolve().parents[4] / "litellm-proxy-extras" @@ -1687,7 +1663,10 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ ): """Day 1 is rolled up and day 2 is still open (never rolled up), so a marker of day 1 must give the same response as reading everything per-key: day 1 from the global table, day 2 - live, one grand total across both. The per-key arm stays on the user table throughout.""" + live, one grand total across both. Per-key rows that land after the rollup then tell the + two sources apart: a late day 1 row is invisible to totals until the next reconcile while a + late day 2 row shows up at once, and both keys rank in the key breakdown, which stays + per-key throughout.""" n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 3 rows: Final = [ ( @@ -1720,39 +1699,56 @@ async def test_get_daily_activity_aggregated_serves_closed_days_from_the_global_ ) _aggregated_postgresql.commit() - async def read(marker: str | None, sql_seen: list[str]): + async def read(marker: str | None): await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) prisma = _prisma_with_marker(marker) - run_query = _psycopg_query_raw(_aggregated_postgresql, []) - - async def query_raw(sql: str, *params: str): - sql_seen.append(sql) - return await run_query(sql, *params) - - prisma.db.query_raw = query_raw + prisma.db.query_raw = _psycopg_query_raw(_aggregated_postgresql, []) return await get_daily_activity_aggregated( prisma_client=prisma, entity_metadata_field=None, **_unfiltered_user_query(), ) - per_key_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim - global_sql: Final[list[str]] = [] # mutable-ok: out-param for the query_raw shim - from_per_key = await read(None, per_key_sql) - from_global = await read("2026-06-01", global_sql) - await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + from_per_key = await read(None) + from_global = await read("2026-06-01") - assert per_key_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 0 - assert global_sql[0].count('FROM "LiteLLM_DailyGlobalSpend"') == 1 - assert global_sql[0].count('FROM "LiteLLM_DailyUserSpend"') == 3 assert from_global.model_dump() == from_per_key.model_dump() - assert from_global.metadata.total_spend == pytest.approx(2 * sum(float(i + 1) for i in range(n_keys))) + seeded_spend: Final = 2 * sum(float(i + 1) for i in range(n_keys)) + assert from_global.metadata.total_spend == pytest.approx(seeded_spend) assert from_global.metadata.total_response_time_ms == 2 * n_keys * 10 * 25 assert from_global.metadata.total_timed_requests == 2 * n_keys assert {day.date.isoformat() for day in from_global.results} == {"2026-06-01", "2026-06-02"} assert len(from_global.results[0].breakdown.api_keys) == USAGE_TOP_API_KEYS_LIMIT assert set(from_global.results[0].breakdown.model_groups) == {"gpt-5", "claude"} + with _aggregated_postgresql.cursor() as cur: + cur.executemany( + """ + INSERT INTO "LiteLLM_DailyUserSpend" + (id, user_id, date, api_key, model, model_group, custom_llm_provider, + endpoint, prompt_tokens, spend, api_requests, successful_requests) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + """, + [ + ("late-1", "user-late", "2026-06-01", "key-late-1", "gpt-5", "", "openai", None, 10, 1000.0, 1, 1), + ("late-2", "user-late", "2026-06-02", "key-late-2", "gpt-5", "", "openai", None, 10, 500.0, 1, 1), + ], + ) + _aggregated_postgresql.commit() + + late_per_key = await read(None) + late_global = await read("2026-06-01") + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + + assert late_per_key.metadata.total_spend == pytest.approx(seeded_spend + 1000.0 + 500.0) + assert late_global.metadata.total_spend == pytest.approx(seeded_spend + 500.0) + by_day: Final = {day.date.isoformat(): day for day in late_global.results} + assert by_day["2026-06-01"].metrics.spend == pytest.approx(seeded_spend / 2) + assert by_day["2026-06-02"].metrics.spend == pytest.approx(seeded_spend / 2 + 500.0) + assert by_day["2026-06-01"].breakdown.api_keys["key-late-1"].metrics.spend == pytest.approx(1000.0) + assert by_day["2026-06-02"].breakdown.api_keys["key-late-2"].metrics.spend == pytest.approx(500.0) + assert late_global.metadata.total_api_keys == n_keys + 2 + @pytest.mark.asyncio async def test_get_daily_activity_aggregated_reports_exact_limit_key_count_as_complete( @@ -1810,7 +1806,20 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo """Rows stored with an empty or NULL model_group must land in the model_groups breakdown under their model name instead of vanishing from the usage UI.""" rows: Final = [ - ("row-0", "user-0", "2026-06-01", "key-0", "gpt-5", "gpt-5-eu", "openai", "/v1/chat/completions", 10, 7.0, 1, 1), + ( + "row-0", + "user-0", + "2026-06-01", + "key-0", + "gpt-5", + "gpt-5-eu", + "openai", + "/v1/chat/completions", + 10, + 7.0, + 1, + 1, + ), ("row-1", "user-1", "2026-06-01", "key-1", "gpt-5", "", "openai", "/v1/chat/completions", 10, 3.0, 1, 1), ("row-2", "user-2", "2026-06-01", "key-2", "claude-x", None, "anthropic", "/v1/messages", 10, 2.0, 1, 1), ] diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index ee72e98ffa9..6121608b658 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -28,6 +28,7 @@ from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch import pytest +from apscheduler.schedulers.asyncio import AsyncIOScheduler from fastapi import FastAPI from pydantic import BaseModel from typing_extensions import TypedDict @@ -1042,8 +1043,8 @@ async def test_spend_report_locks_are_never_released(): proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() -def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, MagicMock]: - scheduler = MagicMock() +def _init_daily_global_spend_reconcile_job() -> tuple[AsyncIOScheduler, MagicMock, MagicMock]: + scheduler = AsyncIOScheduler() proxy_logging_obj = MagicMock() proxy_logging_obj.alerting_handler = AsyncMock() prisma_client = MagicMock() @@ -1058,28 +1059,31 @@ def _init_daily_global_spend_reconcile_job() -> tuple[MagicMock, MagicMock, Magi def test_daily_global_spend_reconcile_job_is_scheduled_nightly_with_an_immediate_catch_up_run(): """Startup schedules the LiteLLM_DailyGlobalSpend backfill a couple of minutes out, so a fresh deploy switches usage reads to the global table without waiting for the nightly - run, and replaces any previous registration of the same job id.""" + run, and after that it fires once a day at 00:30 UTC, when the previous UTC day is closed.""" from datetime import datetime, timedelta, timezone from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID scheduler, _, _ = _init_daily_global_spend_reconcile_job() + job = scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) + assert job is not None - (call,) = scheduler.add_job.call_args_list - assert call.kwargs["id"] == DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID - assert call.kwargs["replace_existing"] is True - assert call.args[1:] == ("cron",) - assert (call.kwargs["hour"], call.kwargs["minute"], call.kwargs["timezone"]) == (0, 30, "UTC") - assert timedelta(0) < call.kwargs["next_run_time"] - datetime.now(timezone.utc) <= timedelta(minutes=2) + assert timedelta(0) < job.next_run_time - datetime.now(timezone.utc) <= timedelta(minutes=2) + after_catch_up = datetime(2026, 9, 16, 12, 0, tzinfo=timezone.utc) + assert job.trigger.get_next_fire_time(None, after_catch_up) == datetime(2026, 9, 17, 0, 30, tzinfo=timezone.utc) + just_after_a_run = datetime(2026, 9, 17, 0, 30, 1, tzinfo=timezone.utc) + assert job.trigger.get_next_fire_time(None, just_after_a_run) == datetime(2026, 9, 18, 0, 30, tzinfo=timezone.utc) @pytest.mark.asyncio async def test_daily_global_spend_reconcile_job_runs_under_the_pod_lock_and_alerts_through_the_proxy(monkeypatch): + from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID + scheduler, proxy_logging_obj, prisma_client = _init_daily_global_spend_reconcile_job() run = AsyncMock() monkeypatch.setattr(ps, "run_scheduled_daily_global_spend_reconcile", run) - await scheduler.add_job.call_args.args[0]() + await scheduler.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func() run.assert_awaited_once() assert run.await_args.args == (prisma_client,) From f25d65940d2a013d1604c729a7dc39836df5e31f Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 12:39:48 +0000 Subject: [PATCH 6/8] fix(proxy): take the closed-day cutoff for the global spend rollup from the database clock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 43 +++++------- .../test_daily_global_spend_rollup.py | 70 +++++++++++-------- 2 files changed, 59 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index ba4a4e4e3d6..c135c7d1d9c 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -12,7 +12,7 @@ a large deployment the first backfill is minutes of work. from collections.abc import Awaitable, Callable from dataclasses import dataclass -from datetime import date, datetime, timedelta, timezone +from datetime import date, timedelta from typing import TYPE_CHECKING, Final from pydantic import BaseModel, ConfigDict, ValidationError @@ -73,7 +73,7 @@ def _reconcile_day_sql() -> str: RECONCILE_DAY_SQL: Final = _reconcile_day_sql() -_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now" +_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now, (NOW() AT TIME ZONE 'UTC')::date::text AS today" _ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' # Pod clocks drift from the database clock and from each other, so rows are picked up from a # little before the previous scan; rewriting a day twice is idempotent. @@ -111,6 +111,7 @@ class _NowRow(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") now: str + today: str @dataclass(frozen=True, slots=True) @@ -160,18 +161,18 @@ async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThroug await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -async def _db_now(prisma_client: "PrismaClient") -> str: +async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) - return _NowRow.model_validate(rows[0]).now + return _NowRow.model_validate(rows[0]) -async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingScan: - """Every closed UTC day (strictly before today) still to roll up, oldest first: days past the - marker, plus any day with per-key rows written since the scan behind the marker. Before a - run has fully succeeded there is no such scan, so every closed day is rolled up.""" +async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan: + """Every closed UTC day (strictly before the database's today) still to roll up, oldest first: + days past the marker, plus any day with per-key rows written since the scan behind the marker. + Before a run has fully succeeded there is no such scan, so every closed day is rolled up.""" marker: Final = await read_marker(prisma_client) - scanned_at: Final = await _db_now(prisma_client) - last_closed_day: Final = (today - timedelta(days=1)).isoformat() + db_now: Final = await _db_now(prisma_client) + last_closed_day: Final = (date.fromisoformat(db_now.today) - timedelta(days=1)).isoformat() rows: Final = ( await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) if marker is None or marker.scanned_at is None @@ -179,11 +180,11 @@ async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingS _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at ) ) - return _PendingScan(marker, scanned_at, tuple(_DateRow.model_validate(row).date for row in rows)) + return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) -async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - return (await _scan_pending(prisma_client, today)).days +async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]: + return (await _scan_pending(prisma_client)).days async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: @@ -192,15 +193,11 @@ async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: await prisma_client.db.execute_raw(RECONCILE_DAY_SQL, day) -async def run_daily_global_spend_reconcile( - prisma_client: "PrismaClient", - today: date | None = None, -) -> ReconcileResult: +async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> ReconcileResult: """Roll up every pending day, advancing the marker after each; a failing day stops the run with the marker on the last good day so the next run resumes there. The scan time is only recorded once every pending day is done, so late rows a failed run saw are found again.""" - effective_today: Final = today or datetime.now(timezone.utc).date() - scan: Final = await _scan_pending(prisma_client, effective_today) + scan: Final = await _scan_pending(prisma_client) done: Final = await _reconcile_until_failure(prisma_client, scan) if len(done) < len(scan.days): marker: Final = await reconciled_through(prisma_client) @@ -243,13 +240,12 @@ async def run_scheduled_daily_global_spend_reconcile( prisma_client: "PrismaClient", pod_lock_manager: "PodLockManager | None" = None, alert: Callable[[str], Awaitable[None]] | None = None, - today: date | None = None, ) -> ReconcileResult | None: """Run the reconcile under a cross-pod lock so one proxy does the work; the lock only saves effort (each day is an idempotent rewrite), so an unreachable Redis runs unguarded rather than skipping.""" redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache if pod_lock_manager is None or redis_cache is None: - return await _run_and_alert(prisma_client, alert=alert, today=today) + return await _run_and_alert(prisma_client, alert=alert) acquired: Final = await pod_lock_manager.acquire_lock( cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID, ttl=DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS @@ -258,7 +254,7 @@ async def run_scheduled_daily_global_spend_reconcile( verbose_proxy_logger.info("Daily global spend reconcile: another pod holds the lock, skipping this run") return None try: - return await _run_and_alert(prisma_client, alert=alert, today=today) + return await _run_and_alert(prisma_client, alert=alert) finally: if acquired: await pod_lock_manager.release_lock(cronjob_id=DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID) @@ -277,9 +273,8 @@ async def _run_and_alert( prisma_client: "PrismaClient", *, alert: Callable[[str], Awaitable[None]] | None, - today: date | None, ) -> ReconcileResult: - result: Final = await run_daily_global_spend_reconcile(prisma_client, today=today) + result: Final = await run_daily_global_spend_reconcile(prisma_client) if result.days_reconciled: verbose_proxy_logger.info( "Daily global spend reconcile: rolled up %d day(s), reconciled through %s", diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index b5b78229c11..9655953134a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -43,7 +43,8 @@ class _FakeConfigTable: class _FakeDb: """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, - so "rows written since the last scan" behaves like Postgres would.""" + so "rows written since the last scan" behaves like Postgres would. The database's own + date decides which day is still open, never the pod's clock.""" def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma @@ -52,7 +53,7 @@ class _FakeDb: async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: if sql.startswith("SELECT (NOW()"): self._prisma.clock += 1 - return [{"now": f"clock-{self._prisma.clock:04d}"}] + return [{"now": f"clock-{self._prisma.clock:04d}", "today": self._prisma.today.isoformat()}] rows = self._prisma.user_rows if len(params) == 1: (last,) = params @@ -73,8 +74,11 @@ class _FakeDb: class _FakePrisma: """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" - def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: + def __init__( + self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY + ) -> None: self.clock = 0 + self.today = today self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days self.reconciled: list[str] = [] @@ -98,12 +102,14 @@ async def _fresh_marker_cache(): @pytest.mark.asyncio -async def test_first_run_rolls_up_every_closed_day_and_never_today(): +async def test_first_run_rolls_up_every_closed_day_and_never_the_database_s_today(): """Before any marker exists every closed day with per-key rows is rolled up. Today is left - out: pods are still flushing it, so it is served live from the per-key table until it closes.""" + out: pods are still flushing it, so it is served live from the per-key table until it closes. + The database clock says which day that is; a pod booting with its clock a day ahead must not + roll the open day up and mark it reconciled.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-03", "2026-09-14", "2026-09-15")) - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-03", "2026-09-14") assert result.failed_day is None @@ -114,11 +120,12 @@ async def test_first_run_rolls_up_every_closed_day_and_never_today(): @pytest.mark.asyncio async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14"), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) prisma.reconciled.clear() + prisma.today = TODAY - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-14",) assert await reconciled_through(prisma) == "2026-09-14" @@ -128,13 +135,14 @@ async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a day far behind the marker. That day is rewritten, and the marker never moves back for it.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13"), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) prisma.reconciled.clear() + prisma.today = TODAY prisma.write_late_row("2026-09-01") prisma.write_late_row("2026-09-03") - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-03") assert "2026-09-05" not in prisma.reconciled @@ -145,15 +153,16 @@ async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_ru async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): """The scan time only advances when every pending day was rewritten, otherwise a late row found by the failed run would be counted as handled.""" - prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13"), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) + prisma.today = TODAY prisma.write_late_row("2026-09-01") prisma.failing_days = frozenset({"2026-09-01"}) - failed = await run_daily_global_spend_reconcile(prisma, today=TODAY) + failed = await run_daily_global_spend_reconcile(prisma) prisma.failing_days = frozenset() prisma.reconciled.clear() - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert failed.failed_day == "2026-09-01" assert failed.reconciled_through == "2026-09-13" @@ -166,7 +175,7 @@ async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-13") marker = await read_marker(prisma) @@ -175,11 +184,11 @@ async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): @pytest.mark.asyncio async def test_a_run_with_no_new_closed_days_keeps_the_marker(): - prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) prisma.reconciled.clear() - result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == () assert result.reconciled_through == "2026-09-13" @@ -191,7 +200,7 @@ async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_goo a global table missing that day's spend.""" prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01",) assert result.failed_day == "2026-09-02" @@ -203,10 +212,10 @@ async def test_a_failing_day_stops_the_run_and_leaves_the_marker_on_the_last_goo @pytest.mark.asyncio async def test_the_next_run_resumes_from_the_failed_day(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-02"})) - await run_daily_global_spend_reconcile(prisma, today=TODAY) + await run_daily_global_spend_reconcile(prisma) prisma.failing_days = frozenset() - result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + result = await run_daily_global_spend_reconcile(prisma) assert result.days_reconciled == ("2026-09-01", "2026-09-02", "2026-09-03") assert await reconciled_through(prisma) == "2026-09-03" @@ -215,13 +224,14 @@ async def test_the_next_run_resumes_from_the_failed_day(): @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" - prisma = _FakePrisma(user_days=("2026-09-13",)) - await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma = _FakePrisma(user_days=("2026-09-13",), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) + prisma.today = TODAY prisma.write_late_row("2026-09-12") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock() - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) assert result is not None assert result.days_reconciled == () @@ -236,7 +246,7 @@ async def test_a_clean_run_does_not_alert(): prisma = _FakePrisma(user_days=("2026-09-13",)) alert = AsyncMock() - await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert, today=TODAY) + await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=None, alert=alert) alert.assert_not_awaited() @@ -256,7 +266,7 @@ async def test_scheduled_run_skips_when_another_pod_holds_the_lock(): prisma = _FakePrisma(user_days=("2026-09-13",)) lock = _pod_lock(acquired=False) - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) assert result is None assert prisma.reconciled == [] @@ -268,7 +278,7 @@ async def test_scheduled_run_runs_and_releases_the_lock_when_it_wins(): prisma = _FakePrisma(user_days=("2026-09-13",)) lock = _pod_lock(acquired=True) - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_awaited_once() @@ -282,7 +292,7 @@ async def test_scheduled_run_proceeds_when_the_lock_cannot_be_acquired_or_read() lock = _pod_lock(acquired=False) lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) - result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock, today=TODAY) + result = await run_scheduled_daily_global_spend_reconcile(prisma, pod_lock_manager=lock) assert result is not None and result.days_reconciled == ("2026-09-13",) lock.release_lock.assert_not_awaited() From abf530fbeb0563d5877ceb883261b25161855b29 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 20:55:11 +0000 Subject: [PATCH 7/8] fix(proxy): never rewind the daily global spend marker from an overlapping reconcile run Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 31 +++++++++++++------ .../test_daily_global_spend_rollup.py | 30 +++++++++++++++++- 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index c135c7d1d9c..d50feb6f16b 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -161,6 +161,24 @@ async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThroug await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) +async def _stored_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: + """The marker as another pod may have just written it, bypassing this pod's config cache.""" + param: Final = await ConfigRepository(prisma_client).get_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + return None if param is None else _marker_from_param_value(param.param_value) + + +async def _record_advanced(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: + """Advance the stored marker by ``days``. Two runs can overlap (Redis unreachable, lock expired + on a long backfill), so the base is what is stored now, not the snapshot this run scanned from: + a slower run may then only add to the faster run's marker, never rewind it. Without a new scan + time the stored one is kept.""" + stored: Final = await _stored_marker(prisma_client) + kept_scanned_at: Final = None if stored is None else stored.scanned_at + await _record_marker( + prisma_client, _advanced(stored, days, scanned_at=scanned_at if scanned_at is not None else kept_scanned_at) + ) + + async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) return _NowRow.model_validate(rows[0]) @@ -203,7 +221,7 @@ async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> Rec marker: Final = await reconciled_through(prisma_client) return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) if scan.marker is not None or done: - await _record_marker(prisma_client, _advanced(scan.marker, done, scanned_at=scan.scanned_at)) + await _record_advanced(prisma_client, done, scanned_at=scan.scanned_at) return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) @@ -215,21 +233,16 @@ def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanne async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: for index, day in enumerate(scan.days): - if not await _reconcile_and_record(prisma_client, scan.marker, scan.days[: index + 1]): + if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): return scan.days[:index] return scan.days -async def _reconcile_and_record( - prisma_client: "PrismaClient", marker: ReconciledThrough | None, done_with_this: tuple[str, ...] -) -> bool: +async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool: day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_marker( - prisma_client, - _advanced(marker, done_with_this, scanned_at=None if marker is None else marker.scanned_at), - ) + await _record_advanced(prisma_client, done_with_this, scanned_at=None) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 9655953134a..69f06fad081 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -40,6 +40,10 @@ class _FakeConfigTable: self.rows[where["param_name"]] = data["update"]["param_value"] return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + async def find_unique(self, *, where: dict[str, str]) -> _FakeConfigRow | None: + stored = self.rows.get(where["param_name"]) + return None if stored is None else _FakeConfigRow(where["param_name"], stored) + class _FakeDb: """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, @@ -68,11 +72,15 @@ class _FakeDb: if day in self._prisma.failing_days: raise RuntimeError(f"day {day} exploded") self._prisma.reconciled.append(day) + landing = self._prisma.marker_landing_on_day.get(day) + if landing is not None: + self.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = landing return 1 class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw. + ``marker_landing_on_day`` stores another pod's marker the moment this run rewrites that day.""" def __init__( self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY @@ -81,6 +89,7 @@ class _FakePrisma: self.today = today self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days + self.marker_landing_on_day: dict[str, str] = {} self.reconciled: list[str] = [] self.db = _FakeDb(self) @@ -221,6 +230,25 @@ async def test_the_next_run_resumes_from_the_failed_day(): assert await reconciled_through(prisma) == "2026-09-03" +@pytest.mark.asyncio +async def test_a_slower_overlapping_run_never_rewinds_the_marker_a_faster_run_stored(): + """Two pods can reconcile at once (Redis unreachable, or the lock expired on a long backfill). + When the faster one has already stored a later marker, the slower one may only add to it. Putting + its own older prefix back, or dropping the scan time, would send usage reads for every day in + between back to the per-key table until the next run.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-03"})) + prisma.marker_landing_on_day = { + "2026-09-02": '{"reconciled_through": "2026-09-14", "scanned_at": "clock-0009"}', + } + + result = await run_daily_global_spend_reconcile(prisma) + + assert result.days_reconciled == ("2026-09-01", "2026-09-02") + assert result.reconciled_through == "2026-09-14" + marker = await read_marker(prisma) + assert marker is not None and (marker.reconciled_through, marker.scanned_at) == ("2026-09-14", "clock-0009") + + @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" From 3449ae9d0d1d339146fa5ffb7318a62553860a46 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 21:15:32 +0000 Subject: [PATCH 8/8] fix(proxy): advance the daily global spend marker in one conditional upsert so overlapping runs cannot rewind it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../daily_global_spend_rollup.py | 51 +++++++--------- .../test_daily_global_spend_rollup.py | 61 ++++++++++++++++--- 2 files changed, 74 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index d50feb6f16b..376b113ed02 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -23,7 +23,6 @@ from litellm.constants import ( DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ) -from litellm.repositories.config_repository import ConfigRepository if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache @@ -82,6 +81,17 @@ _PENDING_DAYS_SQL: Final = ( 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' 'ORDER BY "date"' ) +# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the +# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL. +_ADVANCE_MARKER_SQL: Final = ( + 'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") ' + "VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) " + 'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object(' + "'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', " + "EXCLUDED.\"param_value\" ->> 'reconciled_through'), " + "'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', " + "EXCLUDED.\"param_value\" ->> 'scanned_at'))" +) class ReconciledThrough(BaseModel): @@ -152,33 +162,20 @@ async def reconciled_through(prisma_client: "PrismaClient") -> str | None: return None if marker is None else marker.reconciled_through -async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThrough) -> None: +async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: + """Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later + than what is stored, so a slower overlapping run can only add to a faster run's marker.""" from litellm.proxy.utils import invalidate_config_param - await ConfigRepository(prisma_client).set_param( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, marker.model_dump_json() + await prisma_client.db.execute_raw( + _ADVANCE_MARKER_SQL, + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, + max(days) if days else None, + scanned_at, ) await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -async def _stored_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: - """The marker as another pod may have just written it, bypassing this pod's config cache.""" - param: Final = await ConfigRepository(prisma_client).get_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) - return None if param is None else _marker_from_param_value(param.param_value) - - -async def _record_advanced(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: - """Advance the stored marker by ``days``. Two runs can overlap (Redis unreachable, lock expired - on a long backfill), so the base is what is stored now, not the snapshot this run scanned from: - a slower run may then only add to the faster run's marker, never rewind it. Without a new scan - time the stored one is kept.""" - stored: Final = await _stored_marker(prisma_client) - kept_scanned_at: Final = None if stored is None else stored.scanned_at - await _record_marker( - prisma_client, _advanced(stored, days, scanned_at=scanned_at if scanned_at is not None else kept_scanned_at) - ) - - async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) return _NowRow.model_validate(rows[0]) @@ -221,16 +218,10 @@ async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> Rec marker: Final = await reconciled_through(prisma_client) return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) if scan.marker is not None or done: - await _record_advanced(prisma_client, done, scanned_at=scan.scanned_at) + await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at) return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) -def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanned_at: str | None) -> ReconciledThrough: - """The marker after ``days`` were rewritten: a late old day never moves it back.""" - through: Final = max((marker.reconciled_through if marker is not None else "", *days)) - return ReconciledThrough(reconciled_through=through, scanned_at=scanned_at) - - async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: for index, day in enumerate(scan.days): if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): @@ -242,7 +233,7 @@ async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: t day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_advanced(prisma_client, done_with_this, scanned_at=None) + await _advance_marker(prisma_client, done_with_this, scanned_at=None) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 69f06fad081..3da587435ad 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -1,5 +1,6 @@ """Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" +import json import pathlib import re from datetime import date @@ -14,6 +15,7 @@ from pytest_postgresql import factories from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + _ADVANCE_MARKER_SQL, RECONCILE_DAY_SQL, read_marker, reconciled_through, @@ -36,13 +38,21 @@ class _FakeConfigTable: def __init__(self) -> None: self.rows: dict[str, object] = {} - async def upsert(self, *, where: dict[str, str], data: dict[str, dict[str, str]]) -> _FakeConfigRow: - self.rows[where["param_name"]] = data["update"]["param_value"] - return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + def advance(self, param_name: str, through: str | None, scanned_at: str | None) -> None: + """What ``_ADVANCE_MARKER_SQL`` does in Postgres: keep the later of stored and incoming per field.""" + stored = self.rows.get(param_name) + current: dict[str, str | None] = json.loads(stored) if isinstance(stored, str) else {} + self.rows[param_name] = json.dumps( + { + "reconciled_through": _greatest(current.get("reconciled_through"), through), + "scanned_at": _greatest(current.get("scanned_at"), scanned_at), + } + ) - async def find_unique(self, *, where: dict[str, str]) -> _FakeConfigRow | None: - stored = self.rows.get(where["param_name"]) - return None if stored is None else _FakeConfigRow(where["param_name"], stored) + +def _greatest(stored: str | None, incoming: str | None) -> str | None: + present = [value for value in (stored, incoming) if value is not None] + return max(present) if present else None class _FakeDb: @@ -67,9 +77,14 @@ class _FakeDb: {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) ] - async def execute_raw(self, sql: str, *params: str) -> int: + async def execute_raw(self, sql: str, *params: str | None) -> int: + if sql == _ADVANCE_MARKER_SQL: + param_name, through, scanned_at = params + assert param_name is not None + self.litellm_config.advance(param_name, through, scanned_at) + return 1 (day,) = params - if day in self._prisma.failing_days: + if day is None or day in self._prisma.failing_days: raise RuntimeError(f"day {day} exploded") self._prisma.reconciled.append(day) landing = self._prisma.marker_landing_on_day.get(day) @@ -485,3 +500,33 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")] assert untouched == [] + + +_CONFIG_DDL: Final = 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB)' +_MARKER_SQL: Final = 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s' + + +def test_advance_marker_sql_only_ever_moves_the_stored_marker_forward(_rollup_postgresql: psycopg.Connection): + """Against real Postgres: the statement a slower overlapping run issues after the faster run + already stored a later marker leaves that marker alone, whether it carries an older scan time or + none at all, while a run that is further along moves both fields on.""" + conn: Final = _rollup_postgresql + conn.execute(_CONFIG_DDL) # pyright: ignore[reportArgumentType] # DDL literal + conn.commit() + param: Final = DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM + + def stored() -> object: + with conn.cursor(row_factory=dict_row) as cur: + row = cur.execute(_MARKER_SQL, (param,)).fetchone() + return None if row is None else row["param_value"] + + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-01", None)) + assert stored() == {"reconciled_through": "2026-09-01", "scanned_at": None} + + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-14", "2026-09-15 00:30:02.5")) + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-02", None)) + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-03", "2026-09-15 00:30:01.25")) + assert stored() == {"reconciled_through": "2026-09-14", "scanned_at": "2026-09-15 00:30:02.5"} + + _execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-15", "2026-09-16 00:30:00.75")) + assert stored() == {"reconciled_through": "2026-09-15", "scanned_at": "2026-09-16 00:30:00.75"}