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..d0bc3e159de --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql @@ -0,0 +1,35 @@ +-- 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, + "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, + + 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 cc7793d4bb2..91b59e56906 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -820,6 +820,37 @@ 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) + total_response_time_ms BigInt @default(0) + timed_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 b76f24fc84b..e7cb21a3a7d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -2082,6 +2082,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/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a5e292b177d..a1c92d37871 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -11,6 +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.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, @@ -750,6 +751,66 @@ def _rollup_metric_select(table_name: str) -> str: _MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" +_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", + "total_response_time_ms", + "timed_requests", +) + + +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 + try: + 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 + + +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( *, table_name: str, @@ -762,6 +823,7 @@ 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, + 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. @@ -786,6 +848,7 @@ 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) # TODO: drop the successful_requests/failed_requests aggregates (and the @@ -806,8 +869,7 @@ def _build_aggregated_sql_query( custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level, NULL::bigint AS distinct_api_keys,{metric_select} - FROM "{pg_table}" - WHERE {where_clause} + FROM {_key_free_source(pg_table, where_clause, marker_param)} GROUP BY GROUPING SETS ( (date), (date, model), @@ -850,7 +912,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( @@ -1395,7 +1458,10 @@ 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, + 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 raw_rows, raw_entity_rows = await asyncio.gather( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0652f312114..60f121abe53 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -262,6 +262,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, @@ -688,6 +689,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, @@ -10017,6 +10021,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, @@ -10358,6 +10368,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 cc7793d4bb2..91b59e56906 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -820,6 +820,37 @@ 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) + total_response_time_ms BigInt @default(0) + timed_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..376b113ed02 --- /dev/null +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -0,0 +1,293 @@ +"""Roll closed UTC days of ``LiteLLM_DailyUserSpend`` up into ``LiteLLM_DailyGlobalSpend``. + +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. 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 +from dataclasses import dataclass +from datetime import date, timedelta +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, +) + +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 + +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", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", + "api_requests", + "successful_requests", + "failed_requests", + "total_response_time_ms", + "timed_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: + 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() +_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. +_PENDING_DAYS_SQL: Final = ( + 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ' + '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): + """``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): + 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 + + +class _NowRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + now: str + today: str + + +@dataclass(frozen=True, slots=True) +class ReconcileResult: + days_reconciled: tuple[str, ...] + reconciled_through: str | None + failed_day: str | None = 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: + return ( + ReconciledThrough.model_validate_json(value) + if isinstance(value, str) + else ReconciledThrough.model_validate(value) + ) + except ValidationError: + return None + + +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 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 _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 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 _db_now(prisma_client: "PrismaClient") -> _NowRow: + rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) + return _NowRow.model_validate(rows[0]) + + +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) + 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 + else await prisma_client.db.query_raw( + _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at + ) + ) + return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) + + +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: + """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(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.""" + 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) + return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) + if scan.marker is not None or done: + await _advance_marker(prisma_client, 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", scan: _PendingScan) -> tuple[str, ...]: + for index, day in enumerate(scan.days): + 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", done_with_this: tuple[str, ...]) -> bool: + day: Final = done_with_this[-1] + try: + await reconcile_day(prisma_client, day) + 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 + 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, +) -> 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) + + 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) + 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, +) -> ReconcileResult: + 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", + 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 cc7793d4bb2..91b59e56906 100644 --- a/schema.prisma +++ b/schema.prisma @@ -820,6 +820,37 @@ 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) + total_response_time_ms BigInt @default(0) + timed_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/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index e7f8bcc7e4f..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 @@ -1,3 +1,4 @@ +import pathlib import re from collections.abc import Sequence from datetime import datetime, timedelta, timezone @@ -10,10 +11,11 @@ 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 - - -from litellm.constants import PTU_SENTINEL_API_KEY, USAGE_TOP_API_KEYS_LIMIT +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 +25,12 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, + 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, SpendMetrics, @@ -1578,6 +1584,172 @@ 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", {}, "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), + ("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_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 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_global_rollup_marker_read_failure_falls_back_to_the_per_key_table(): + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + prisma = _prisma_with_marker(None) + prisma.get_generic_data = AsyncMock(side_effect=RuntimeError("db down")) + + assert await global_rollup_reconciled_through(prisma, _unfiltered_user_query()) is None + 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_serves_closed_days_from_the_global_table_and_open_days_live( + _aggregated_postgresql: psycopg.Connection, +): + """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. 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 = [ + ( + 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( + '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 + {"p1": "2026-06-01"}, + ) + _aggregated_postgresql.commit() + + async def read(marker: str | None): + await evict_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + prisma = _prisma_with_marker(marker) + 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(), + ) + + from_per_key = await read(None) + from_global = await read("2026-06-01") + + assert from_global.model_dump() == from_per_key.model_dump() + 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( _aggregated_postgresql: psycopg.Connection, @@ -1634,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 deb7289d2d1..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,6 +1043,57 @@ 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[AsyncIOScheduler, MagicMock, MagicMock]: + scheduler = AsyncIOScheduler() + 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 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 + + 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.get_job(DAILY_GLOBAL_SPEND_RECONCILE_JOB_ID).func() + + 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..3da587435ad --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -0,0 +1,532 @@ +"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818).""" + +import json +import pathlib +import re +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, merge_by_conflict_key +from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( + _ADVANCE_MARKER_SQL, + RECONCILE_DAY_SQL, + read_marker, + 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] = {} + + 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), + } + ) + + +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: + """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. 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 + self.litellm_config = _FakeConfigTable() + + 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}", "today": self._prisma.today.isoformat()}] + 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 | 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 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) + 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. + ``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 + ) -> 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.marker_landing_on_day: dict[str, str] = {} + 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) + + +@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_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. + 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) + + 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-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_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"), 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) + + 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"), 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) + + 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"), 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) + prisma.failing_days = frozenset() + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma) + + 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) + + 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",), today=date(2026, 9, 14)) + await run_daily_global_spend_reconcile(prisma) + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma) + + assert result.days_reconciled == () + assert result.reconciled_through == "2026-09-13" + + +@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) + + 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) + prisma.failing_days = frozenset() + + 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" + + +@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.""" + 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) + + 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) + + 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) + + 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) + + assert result is not None and result.days_reconciled == ("2026-09-13",) + 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) + + assert result is not None and result.days_reconciled == ("2026-09-13",) + 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, + 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) + ) +""" + +_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(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, + total_response_time_ms, timed_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, + "total_response_time_ms": 800, + "timed_requests": 1, + **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"]), + int(r["total_response_time_ms"]), + int(r["timed_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: 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 + 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(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 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"}