diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql new file mode 100644 index 00000000000..81d32345794 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at" +ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/litellm/constants.py b/litellm/constants.py index 665ab771726..271e42ee661 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1895,6 +1895,7 @@ _background_interaction_cost_polling_env: Final = os.getenv( "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true" ).lower() BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true" +BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0 PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 21d438d9621..275eb17787e 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -32,19 +32,24 @@ task to settle. A ``BackgroundSettlementStore`` makes the pending settlement durable across processes: the create registers the request context that billing needs (never provider credentials), the settlement is claimed exactly once through the store, and a delete on any replica rebuilds the -billing context from the store when the poll task is not local. Rows left -unclaimed by a process that died are resumed at startup. The default store is -in-memory, which keeps the SDK and single-process behavior unchanged; the -proxy installs a database-backed one. +billing context from the store when the poll task is not local. Every row +still unsettled when a process starts is resumed there, including one whose +settler died between its claim and its outcome: the database-backed store +treats a claim as a lease that its settler renews until it records the +outcome, so such a row is claimable again once the lease runs out. The +default store is in-memory, which keeps the SDK and single-process behavior +unchanged; the proxy installs a database-backed one. """ import asyncio +import itertools from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias +from openai import APIStatusError from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from pydantic_core import PydanticSerializationError, to_jsonable_python @@ -54,10 +59,12 @@ from litellm.constants import ( BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS, BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS, BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.router import CredentialLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams if TYPE_CHECKING: @@ -101,6 +108,7 @@ class PendingBackgroundInteraction: custom_llm_provider: str create_context: BackgroundInteractionCreateContext created_at: datetime + claim_expires_at: datetime | None = None class BackgroundSettlementStore(Protocol): @@ -112,9 +120,11 @@ class BackgroundSettlementStore(Protocol): async def claim(self, interaction_id: str) -> bool: ... + async def renew_claim(self, interaction_id: str) -> None: ... + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ... - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: ... + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: ... @dataclass(frozen=True, slots=True) @@ -145,10 +155,13 @@ class InMemoryBackgroundSettlementStore: self._rows[interaction_id] = None return True + async def renew_claim(self, interaction_id: str) -> None: + return None + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: return None - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: return tuple(row for row in self._rows.values() if row is not None) @@ -169,11 +182,15 @@ class BackgroundInteractionPollContext: interaction_id: str custom_llm_provider: str logging_obj: "LiteLLMLoggingObj" - api_key: str | None = None - api_base: str | None = None + credentials: CredentialLiteLLMParams = field(default_factory=CredentialLiteLLMParams) + deployment_id: str | None = None + model_group: str | None = None initial_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS + start_delay_seconds: float = 0.0 + outcome_write_timeout_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 2 + claim_renewal_interval_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 3 store: BackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore) resumed: bool = False @@ -184,15 +201,41 @@ FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Await async def fetch_background_interaction(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: from litellm.interactions import aget + credentials: Final = context.credentials return await aget( interaction_id=context.interaction_id, custom_llm_provider=context.custom_llm_provider, - api_key=context.api_key, - api_base=context.api_base, + api_key=credentials.api_key, + api_base=credentials.api_base, + vertex_project=credentials.vertex_project, + vertex_location=credentials.vertex_location, + vertex_credentials=credentials.vertex_credentials, **{"no-log": True}, ) +_FETCH_CREDENTIAL_FIELDS: Final = ("api_key", "api_base", "vertex_project", "vertex_location", "vertex_credentials") + + +def _fetch_credentials(request_kwargs: Mapping[str, object]) -> CredentialLiteLLMParams: + try: + return CredentialLiteLLMParams.model_validate( + {field_name: request_kwargs.get(field_name) for field_name in _FETCH_CREDENTIAL_FIELDS} + ) + except ValidationError: + verbose_logger.debug( + "Fetching a background interaction with environment credentials: the request's are malformed" + ) + return CredentialLiteLLMParams() + + +_PERMANENT_FETCH_STATUS_CODES: Final = frozenset({400, 401, 403, 404}) + + +def _is_permanent_fetch_error(error: Exception) -> bool: + return isinstance(error, APIStatusError) and error.status_code in _PERMANENT_FETCH_STATUS_CODES + + def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]: elapsed = 0.0 # rebind-ok: the schedule accumulates the time it has already yielded interval = initial # rebind-ok: the schedule doubles the interval up to the cap @@ -231,6 +274,7 @@ _UNCARRIED_METADATA_KEY: Final = "user_api_key_auth" _JSON_VALUE: Final = TypeAdapter(JsonValue) _STRING: Final = TypeAdapter(str) _OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) def _carries(key: str) -> bool: @@ -256,10 +300,36 @@ def _as_datetime(start_time: datetime | float) -> datetime: return start_time if isinstance(start_time, datetime) else datetime.fromtimestamp(start_time, tz=timezone.utc) +def _spend_metadata(logging_obj: "LiteLLMLoggingObj") -> Mapping[str, object]: + """ + The metadata the creating worker's own spend row is built from, with the + request tags computed once, header-derived ones included, because a + rebuilt logging object has no request headers to derive them from. + """ + litellm_params: Final = _OBJECT_DICT.validate_python(logging_obj.model_call_details.get("litellm_params") or {}) + return { + **_OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)), + **_request_tags(litellm_params), + } + + +def _request_tags(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + try: + return { + "tags": StandardLoggingPayloadSetup.get_request_tags( + litellm_params=_OBJECT_DICT.validate_python(litellm_params), + proxy_server_request=_OBJECT_DICT.validate_python(litellm_params.get("proxy_server_request") or {}), + ) + } + except AttributeError: + verbose_logger.debug("Keeping a background interaction's request tags as sent: they are not a list") + return {} + + def _create_context(logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str) -> BackgroundInteractionCreateContext: - metadata: Final = _OBJECT_MAPPING.validate_python( - get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) - ) + metadata: Final = _spend_metadata(logging_obj) litellm_params: Final = _OBJECT_MAPPING.validate_python(logging_obj.litellm_params) model: Final = logging_obj.model_call_details.get("model") return BackgroundInteractionCreateContext( @@ -313,6 +383,16 @@ async def _settled_elsewhere(context: BackgroundInteractionPollContext) -> bool: return False +async def _settled_elsewhere_before_the_window_ends( + context: BackgroundInteractionPollContext, intervals: Iterator[float] +) -> bool: + for interval in intervals: + await asyncio.sleep(interval) + if await _settled_elsewhere(context): + return True + return False + + async def _claim(context: BackgroundInteractionPollContext) -> bool | None: """ Exactly-once gate between every settler of one interaction, on every @@ -326,30 +406,118 @@ async def _claim(context: BackgroundInteractionPollContext) -> bool | None: return None -async def _record(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> SettlementOutcome: +async def _wrote_outcome(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool: try: await context.store.record_outcome(context.interaction_id, outcome) - except Exception: # noqa: BLE001 # the outcome is an audit trail; the claim already made the settlement exclusive - verbose_logger.exception("Could not record the settlement of background interaction %s", context.interaction_id) + except Exception as e: # noqa: BLE001 # a failed write is retried while the claim still shuts out takeovers + verbose_logger.warning( + "Could not record the settlement of background interaction %s: %s", context.interaction_id, e + ) + return False + return True + + +async def _recorded(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool: + """ + Retries a failed write for half the claim lease, while no other settler + can take the row over, so a database blip does not leave the row to a + takeover or hold the budget reservation that only a recorded release frees. + """ + retries: Final = _poll_intervals( + initial=context.initial_interval_seconds, + maximum=context.max_interval_seconds, + timeout=context.outcome_write_timeout_seconds, + ) + for delay in itertools.chain((0.0,), retries): + await asyncio.sleep(delay) + if await _wrote_outcome(context, outcome): + return True + verbose_logger.error( + "Could not record the settlement of background interaction %s within %ss; the next process to take " + "over its lapsed claim settles it", + context.interaction_id, + context.outcome_write_timeout_seconds, + ) + return False + + +async def _bill(context: BackgroundInteractionPollContext, response: InteractionsAPIResponse) -> SettlementOutcome: + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) + await _recorded(context, "billed") + return "billed" + + +async def _release( + context: BackgroundInteractionPollContext, outcome: Literal["released", "unsettled"] +) -> SettlementOutcome: + """ + A release writes no spend row, so a takeover cannot tell whether the + settler before it released. The outcome is recorded first and the + reservation released only once it is, which leaves a takeover with + nothing to release twice. + """ + if await _recorded(context, outcome): + await _release_open_budget_reservation(logging_obj=context.logging_obj) return outcome +async def _renew_claim_until_cancelled(context: BackgroundInteractionPollContext) -> None: + while True: + await asyncio.sleep(context.claim_renewal_interval_seconds) + try: + await context.store.renew_claim(context.interaction_id) + except Exception as e: # noqa: BLE001 # a missed renewal is retried next interval, well inside the lease + verbose_logger.warning( + "Could not renew the settlement claim of background interaction %s: %s", context.interaction_id, e + ) + + +async def _settle( + context: BackgroundInteractionPollContext, response: InteractionsAPIResponse | None +) -> SettlementOutcome: + """ + Billing runs every success callback in turn, and a slow one can outlast + the claim lease, so the claim is renewed until the outcome is recorded. + A settler that stops renewing has died, and only then can another take + the row over. + """ + renewal: Final = asyncio.create_task(_renew_claim_until_cancelled(context)) + try: + if response is not None and response.status in _TERMINAL_STATUSES and response.usage is not None: + return await _bill(context, response) + return await _release(context, "released") + finally: + renewal.cancel() + + async def poll_and_log_background_interaction_cost( context: BackgroundInteractionPollContext, fetch_interaction: FetchInteraction = fetch_background_interaction, ) -> SettlementOutcome | None: last_response: InteractionsAPIResponse | None = None # rebind-ok: the give-up path settles from the last poll - for interval in _poll_intervals( + await asyncio.sleep(context.start_delay_seconds) + intervals: Final = _poll_intervals( initial=context.initial_interval_seconds, maximum=context.max_interval_seconds, timeout=context.timeout_seconds, - ): + ) + for interval in intervals: await asyncio.sleep(interval) if await _settled_elsewhere(context): return None try: response = await fetch_interaction(context) except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop + if context.resumed and _is_permanent_fetch_error(e): + verbose_logger.warning( + "Stopped fetching resumed background interaction %s: the provider answered %s, so this " + "process cannot read it and leaves it to the replica that created it until the poll window ends", + context.interaction_id, + e, + ) + if await _settled_elsewhere_before_the_window_ends(context, intervals): + return None + break verbose_logger.debug( "Background interaction cost poll for %s failed, will retry: %s", context.interaction_id, @@ -363,11 +531,11 @@ async def poll_and_log_background_interaction_cost( continue if not claimed: return None - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response)) + return await _settle(context, response) if not await _claim(context): return None if last_response is not None and last_response.status in _TERMINAL_STATUSES: - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=last_response)) + return await _settle(context, last_response) if last_response is not None and last_response.status not in _POLLABLE_STATUSES: verbose_logger.error( "Gave up cost polling for background interaction %s after %ss: its last status %r is in neither " @@ -379,20 +547,10 @@ async def poll_and_log_background_interaction_cost( ) else: verbose_logger.warning( - "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", + "Gave up cost polling for background interaction %s; its usage will not be tracked", context.interaction_id, - context.timeout_seconds, ) - await _release_open_budget_reservation(logging_obj=context.logging_obj) - return await _record(context, "unsettled") - - -async def _settle_terminal(logging_obj: "LiteLLMLoggingObj", response: InteractionsAPIResponse) -> SettlementOutcome: - if response.status in _TERMINAL_STATUSES and response.usage is not None: - await _bill_settled_interaction(logging_obj=logging_obj, response=response) - return "billed" - await _release_open_budget_reservation(logging_obj=logging_obj) - return "released" + return await _release(context, "unsettled") async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None: @@ -405,7 +563,9 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> settlement must release the reservation here or the spend counters stay pinned at the estimated cost. """ - metadata: Final = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + metadata: Final = _OBJECT_MAPPING.validate_python( + get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + ) budget_reservation: Final = metadata.get("user_api_key_budget_reservation") if not isinstance(budget_reservation, dict): return @@ -518,13 +678,18 @@ class _UnverifiedRegistrationStore: return False return await self.local.claim(interaction_id) + async def renew_claim(self, interaction_id: str) -> None: + if await self.local.is_claimed(interaction_id): + return + await self.durable.renew_claim(interaction_id) + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: if await self.local.is_claimed(interaction_id): return await self.durable.record_outcome(interaction_id, outcome) - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: - return await self.durable.unclaimed() + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: + return await self.durable.unsettled() async def _registered_store( @@ -559,8 +724,6 @@ async def maybe_schedule_background_interaction_cost_polling( logging_obj: Final = create_kwargs.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return None - api_key: Final = create_kwargs.get("api_key") - api_base: Final = create_kwargs.get("api_base") pending: Final = PendingBackgroundInteraction( interaction_id=response.id, custom_llm_provider=custom_llm_provider, @@ -571,8 +734,7 @@ async def maybe_schedule_background_interaction_cost_polling( interaction_id=response.id, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, - api_key=api_key if isinstance(api_key, str) else None, - api_base=api_base if isinstance(api_base, str) else None, + credentials=_fetch_credentials(create_kwargs), store=await _registered_store(store or _STORE.store, pending), ) return _track_poll(context, fetch_interaction) @@ -603,10 +765,7 @@ async def _settle_before_delete( ) -> SettlementOutcome | None: if not await _claim(context): return None - if response is None: - await _release_open_budget_reservation(logging_obj=context.logging_obj) - return await _record(context, "released") - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response)) + return await _settle(context, response) async def maybe_settle_background_interaction_before_delete( @@ -622,14 +781,11 @@ async def maybe_settle_background_interaction_before_delete( pending: Final = await _pending(settlement_store, interaction_id) if pending is None: return None - api_key: Final = delete_kwargs.get("api_key") - api_base: Final = delete_kwargs.get("api_base") context: Final = BackgroundInteractionPollContext( interaction_id=interaction_id, custom_llm_provider=pending.custom_llm_provider, logging_obj=_rebuild_logging_obj(pending.create_context), - api_key=api_key if isinstance(api_key, str) else None, - api_base=api_base if isinstance(api_base, str) else None, + credentials=_fetch_credentials(delete_kwargs), store=settlement_store, ) try: @@ -644,9 +800,9 @@ async def maybe_settle_background_interaction_before_delete( return await _settle_before_delete(context, response) -async def _unclaimed(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]: +async def _unsettled(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]: try: - return await store.unclaimed() + return await store.unsettled() except Exception: # noqa: BLE001 # an unreadable store at startup leaves its rows for the next boot verbose_logger.exception("Could not list the unsettled background interactions") return () @@ -662,17 +818,36 @@ class PollSchedule: DEFAULT_POLL_SCHEDULE: Final = PollSchedule() +def _deployment_id(create_context: BackgroundInteractionCreateContext) -> str | None: + model_info: Final = create_context.metadata.get("model_info") + deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None + return deployment_id if isinstance(deployment_id, str) else None + + +def _model_group(create_context: BackgroundInteractionCreateContext) -> str | None: + model_group: Final = create_context.metadata.get("model_group") + return model_group if isinstance(model_group, str) else None + + +def _seconds_until(moment: datetime | None) -> float: + if moment is None: + return 0.0 + return max((moment - datetime.now(timezone.utc)).total_seconds(), 0.0) + + def _resumed_context( row: PendingBackgroundInteraction, store: BackgroundSettlementStore, schedule: PollSchedule ) -> BackgroundInteractionPollContext: - age_seconds: Final = (datetime.now(timezone.utc) - row.created_at).total_seconds() return BackgroundInteractionPollContext( interaction_id=row.interaction_id, custom_llm_provider=row.custom_llm_provider, logging_obj=_rebuild_logging_obj(row.create_context), + deployment_id=_deployment_id(row.create_context), + model_group=_model_group(row.create_context), initial_interval_seconds=schedule.initial_interval_seconds, max_interval_seconds=schedule.max_interval_seconds, - timeout_seconds=max(schedule.timeout_seconds - age_seconds, schedule.initial_interval_seconds), + timeout_seconds=schedule.timeout_seconds, + start_delay_seconds=_seconds_until(row.claim_expires_at), store=store, resumed=True, ) @@ -684,13 +859,15 @@ async def resume_unsettled_background_interactions( schedule: PollSchedule = DEFAULT_POLL_SCHEDULE, ) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]: """ - Pick up every settlement no process has claimed, which is what a replica - that died mid-poll leaves behind. Each resumed poll keeps the remaining - share of the original timeout and gets at least one fetch, so a completed - interaction is still billed however late the resume comes. + Pick up every settlement no process has finished, which is what a replica + that died mid-poll or mid-settlement leaves behind. A row whose claim is + still within its lease waits until the lease runs out, then polls like any + other. Each resumed poll gets the full poll window however old its row is, + so a fetch that fails right after a late restart is retried rather than + ending the row as unsettled. """ return tuple( _track_poll(_resumed_context(row, store, schedule), fetch_interaction) - for row in await _unclaimed(store) + for row in await _unsettled(store) if row.interaction_id not in _ACTIVE_POLLS ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 58037689597..9eb2bc93391 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1524,7 +1524,6 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState await asyncio.sleep(5) asyncio.create_task(_run_agent_grant_id_migration()) - await install_background_interaction_settlement(prisma_client) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1651,6 +1650,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ## SYNC UI SETTINGS ## await ProxyStartupEvent.sync_ui_settings_to_general_settings() + await install_background_interaction_settlement( + prisma_client, router=lambda: llm_router, spend_logs_kept=lambda: disable_spend_logs is False + ) + # Start background health checks AFTER models are loaded and index is built if use_background_health_checks: asyncio.create_task(_run_background_health_check()) # start the background health check coroutine. diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/litellm/proxy/spend_tracking/background_interaction_settlement.py b/litellm/proxy/spend_tracking/background_interaction_settlement.py index 165f813cd50..ffa8185619c 100644 --- a/litellm/proxy/spend_tracking/background_interaction_settlement.py +++ b/litellm/proxy/spend_tracking/background_interaction_settlement.py @@ -1,21 +1,25 @@ import asyncio import os import socket -from collections.abc import Awaitable, Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime, timezone +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeVar from pydantic import ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED +from litellm.constants import ( + BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, +) from litellm.interactions.background_cost_polling import ( DEFAULT_POLL_SCHEDULE, BackgroundInteractionCreateContext, + BackgroundInteractionPollContext, FetchInteraction, PendingBackgroundInteraction, PollSchedule, @@ -24,10 +28,13 @@ from litellm.interactions.background_cost_polling import ( fetch_background_interaction, resume_unsettled_background_interactions, ) -from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository +from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository, SpendLogsRepository +from litellm.types.interactions import InteractionsAPIResponse +from litellm.types.router import CredentialLiteLLMParams if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient + from litellm.router import Router class _SettlementRow(Protocol): @@ -41,6 +48,8 @@ class _SettlementRow(Protocol): def created_at(self) -> datetime: ... @property def claimed_at(self) -> datetime | None: ... + @property + def settled_at(self) -> datetime | None: ... class _NewSettlementRow(TypedDict): @@ -59,8 +68,24 @@ class _UnclaimedRowKey(TypedDict): claimed_at: ReadOnly[None] -class _UnclaimedRows(TypedDict): - claimed_at: ReadOnly[None] +class _Before(TypedDict): + lt: ReadOnly[datetime] + + +class _ExpiredClaimRowKey(TypedDict): + interaction_id: ReadOnly[str] + settled_at: ReadOnly[None] + claimed_at: ReadOnly[_Before] + + +class _OwnClaimRowKey(TypedDict): + interaction_id: ReadOnly[str] + claimed_by: ReadOnly[str] + settled_at: ReadOnly[None] + + +class _UnsettledRows(TypedDict): + settled_at: ReadOnly[None] class _Claim(TypedDict): @@ -79,9 +104,19 @@ class _SettlementTableActions(Protocol): def find_unique(self, *, where: _RowKey) -> Awaitable[_SettlementRow | None]: ... - def find_many(self, *, where: _UnclaimedRows) -> Awaitable[Sequence[_SettlementRow]]: ... + def find_many(self, *, where: _UnsettledRows) -> Awaitable[Sequence[_SettlementRow]]: ... - def update_many(self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey) -> Awaitable[int]: ... + def update_many( + self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey | _ExpiredClaimRowKey | _OwnClaimRowKey + ) -> Awaitable[int]: ... + + +class _SpendLogKey(TypedDict): + request_id: ReadOnly[str] + + +class _SpendLogLookup(Protocol): + def find_unique(self, *, where: _SpendLogKey) -> Awaitable[object | None]: ... _CLEARED_CREATE_CONTEXT: Final[Mapping[str, object]] = MappingProxyType({}) @@ -103,11 +138,7 @@ def _json(data: Mapping[str, object]) -> object: return Json.keys(**data) -def _pending_rows(rows: Sequence[_SettlementRow]) -> tuple[PendingBackgroundInteraction, ...]: - return tuple(chain.from_iterable(_pending_row(row) for row in rows)) - - -def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ...]: +def _pending_row(row: _SettlementRow, claim_expires_at: datetime | None) -> tuple[PendingBackgroundInteraction, ...]: try: create_context: Final = BackgroundInteractionCreateContext.model_validate(row.create_context) except ValidationError: @@ -122,14 +153,31 @@ def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ... custom_llm_provider=row.custom_llm_provider, create_context=create_context, created_at=row.created_at, + claim_expires_at=claim_expires_at, ), ) @dataclass(frozen=True, slots=True) class PrismaBackgroundSettlementStore: + """ + A claim is a lease, not a settlement: the settler that claims a row bills + or releases it and then records the outcome, so a row claimed longer ago + than the lease with no outcome belongs to a settler that died in between, + and any settler may claim it again. A live settler renews its claim while + it settles, so it is never overtaken. The settler before may have billed + and then died or failed to record it, so a takeover bills only when the + interaction has no spend row yet, and otherwise records it as billed. The + spend row is read before the claim is taken, so a failed read takes + nothing and the claim is retried. With spend logs off there is no row to + tell, so a takeover bills nothing rather than risk billing twice. + """ + table: _SettlementTableActions + spend_logs: _SpendLogLookup + spend_logs_kept: Callable[[], bool] claimed_by: str + claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS async def register(self, pending: PendingBackgroundInteraction) -> None: await self.table.create( @@ -143,24 +191,71 @@ class PrismaBackgroundSettlementStore: async def pending(self, interaction_id: str) -> PendingBackgroundInteraction | None: row: Final = await self._row(interaction_id) - if row is None or row.claimed_at is not None: + if row is None or self._is_owned(row): return None - return next(iter(_pending_row(row)), None) + return next(iter(_pending_row(row, claim_expires_at=None)), None) async def is_claimed(self, interaction_id: str) -> bool: row: Final = await self._row(interaction_id) - return row is not None and row.claimed_at is not None + return row is not None and self._is_owned(row) async def claim(self, interaction_id: str) -> bool: + if await self._claim_where(_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None)): + return True + settled_before: Final = await self._outcome_a_dead_settler_left(interaction_id) + took_over: Final = await self._claim_where( + _ExpiredClaimRowKey( + interaction_id=interaction_id, + settled_at=None, + claimed_at=_Before(lt=datetime.now(timezone.utc) - self._lease), + ) + ) + if not took_over or settled_before is None: + return took_over + if settled_before == "unsettled": + verbose_proxy_logger.warning( + "Background interaction %s lost its settler mid-settle, and with spend logs off nothing shows " + "whether it was billed, so it is left unbilled rather than risk billing it twice", + interaction_id, + ) + await self.record_outcome(interaction_id, settled_before) + return False + + async def renew_claim(self, interaction_id: str) -> None: + await self.table.update_many( + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), + where=_OwnClaimRowKey(interaction_id=interaction_id, claimed_by=self.claimed_by, settled_at=None), + ) + + async def _claim_where(self, where: _UnclaimedRowKey | _ExpiredClaimRowKey) -> bool: claimed_rows: Final = await _read_from_a_table_that_may_not_exist( self.table.update_many( - data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), - where=_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None), + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), where=where ), when_missing=0, ) return claimed_rows == 1 + async def _outcome_a_dead_settler_left(self, interaction_id: str) -> Literal["billed", "unsettled"] | None: + if not self.spend_logs_kept(): + return "unsettled" + if await self.spend_logs.find_unique(where=_SpendLogKey(request_id=interaction_id)) is None: + return None + return "billed" + + @property + def _lease(self) -> timedelta: + return timedelta(seconds=self.claim_lease_seconds) + + def _claim_expires_at(self, row: _SettlementRow) -> datetime | None: + return None if row.claimed_at is None else row.claimed_at + self._lease + + def _is_owned(self, row: _SettlementRow) -> bool: + if row.settled_at is not None: + return True + claim_expires_at: Final = self._claim_expires_at(row) + return claim_expires_at is not None and claim_expires_at > datetime.now(timezone.utc) + async def _row(self, interaction_id: str) -> _SettlementRow | None: return await _read_from_a_table_that_may_not_exist( self.table.find_unique(where=_RowKey(interaction_id=interaction_id)), when_missing=None @@ -174,19 +269,78 @@ class PrismaBackgroundSettlementStore: where=_RowKey(interaction_id=interaction_id), ) - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: - return _pending_rows(await self.table.find_many(where=_UnclaimedRows(claimed_at=None))) + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: + rows: Final = await self.table.find_many(where=_UnsettledRows(settled_at=None)) + return tuple(chain.from_iterable(_pending_row(row, self._claim_expires_at(row)) for row in rows)) + + +def _provider_credentials(router: "Router", deployment_id: str, custom_llm_provider: str) -> Mapping[str, object]: + credentials: Final = router.get_deployment_credentials_with_provider(model_id=deployment_id) + if credentials is None or credentials.get("custom_llm_provider") != custom_llm_provider: + return MappingProxyType({}) + return MappingProxyType(credentials) + + +def _group_deployment_credentials(router: "Router", context: BackgroundInteractionPollContext) -> Mapping[str, object]: + if context.model_group is None: + return MappingProxyType({}) + candidates: Final = tuple( + credentials + for deployment_id in router.get_model_ids(model_name=context.model_group) + if (credentials := _provider_credentials(router, deployment_id, context.custom_llm_provider)) + ) + return candidates[0] if len(candidates) == 1 else MappingProxyType({}) + + +def _deployment_credentials( + router: "Router | None", context: BackgroundInteractionPollContext +) -> CredentialLiteLLMParams | None: + if router is None: + return None + serving: Final = ( + _provider_credentials(router, context.deployment_id, context.custom_llm_provider) + if context.deployment_id is not None and router.get_deployment(model_id=context.deployment_id) is not None + else _group_deployment_credentials(router, context) + ) + return CredentialLiteLLMParams.model_validate(serving) if serving else None + + +def fetch_with_deployment_credentials( + router: Callable[[], "Router | None"], fetch_interaction: FetchInteraction = fetch_background_interaction +) -> FetchInteraction: + """ + A resumed poll has no request to take credentials from, and they are never + stored, so it fetches with the credentials of the deployment that served + the create, looked up by its id on the live router at each fetch so a + deployment loaded from the database after startup is found too. Editing a + deployment changes its id, so a missing id falls back to the create's model + group when exactly one of its deployments serves the same provider. Anything + else falls back to the provider's environment credentials, never to a + wildcard route, whose key belongs to whatever provider it routes to. + """ + + async def fetch(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: + credentials: Final = _deployment_credentials(router(), context) + if credentials is None: + return await fetch_interaction(context) + return await fetch_interaction(replace(context, credentials=credentials)) + + return fetch async def configure_background_interaction_settlement( table: _SettlementTableActions, + spend_logs: _SpendLogLookup, + spend_logs_kept: Callable[[], bool], claimed_by: str, fetch_interaction: FetchInteraction = fetch_background_interaction, schedule: PollSchedule = DEFAULT_POLL_SCHEDULE, ) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]: if not BACKGROUND_INTERACTION_COST_POLLING_ENABLED: return () - store: Final = PrismaBackgroundSettlementStore(table=table, claimed_by=claimed_by) + store: Final = PrismaBackgroundSettlementStore( + table=table, spend_logs=spend_logs, spend_logs_kept=spend_logs_kept, claimed_by=claimed_by + ) configure_background_settlement_store(store) resumed: Final = await resume_unsettled_background_interactions(store, fetch_interaction, schedule) if resumed: @@ -194,11 +348,16 @@ async def configure_background_interaction_settlement( return resumed -async def install_background_interaction_settlement(prisma_client: "PrismaClient") -> None: +async def install_background_interaction_settlement( + prisma_client: "PrismaClient", router: Callable[[], "Router | None"], spend_logs_kept: Callable[[], bool] +) -> None: try: await configure_background_interaction_settlement( table=BackgroundInteractionSettlementRepository(prisma_client).table, + spend_logs=SpendLogsRepository(prisma_client).table, + spend_logs_kept=spend_logs_kept, claimed_by=f"{socket.gethostname()}:{os.getpid()}", + fetch_interaction=fetch_with_deployment_credentials(router), ) except Exception as e: # noqa: BLE001 # a boot step must survive any DB error; billing then settles in-process as before verbose_proxy_logger.warning( diff --git a/schema.prisma b/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 16eb0a3434e..4408bd147f9 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -1,3 +1,4 @@ +import json import math import socket import time @@ -17,6 +18,8 @@ from integration._support.client import ( Gateway, eventually, gateway_from_environment, + held, + list_value, object_value, string_value, ) @@ -54,6 +57,9 @@ _CUSTOM_INPUT_RATE: Final = 2e-06 _CUSTOM_OUTPUT_RATE: Final = 4e-05 _ENV_KEY: Final = "integration-gemini-env-key" _DEPLOYMENT_KEY: Final = "integration-gemini-deployment-key" +_DATABASE_DEPLOYMENT_KEY: Final = "integration-gemini-database-deployment-key" +_ROTATED_DEPLOYMENT_KEY: Final = "integration-gemini-rotated-deployment-key" +_TAG: Final = "integration-settlement-team" _CREATOR_POLL: Final = {"BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "300"} _SETTLER_POLL: Final = { "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "1", @@ -65,8 +71,10 @@ _RESUMER_POLL: Final = { "BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS": "1", "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "120", } +_EAGER_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "0.01"} +_LATE_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "30"} _SPEND_QUERY: Final = ( - "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens " + "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens, request_tags " 'FROM "LiteLLM_SpendLogs" WHERE request_id = %s' ) _SETTLEMENT_QUERY: Final = ( @@ -74,6 +82,16 @@ _SETTLEMENT_QUERY: Final = ( 'settled_at IS NOT NULL AS settled, create_context FROM "LiteLLM_BackgroundInteractionSettlement" ' "WHERE interaction_id = %s" ) +_ABANDONED_CLAIM: Final = ( + 'UPDATE "LiteLLM_BackgroundInteractionSettlement" SET claimed_by = %s, ' + "claimed_at = (now() AT TIME ZONE 'UTC') - interval '10 minutes', settled_at = NULL, outcome = NULL, " + "create_context = %s::jsonb WHERE interaction_id = %s" +) +_DEAD_SETTLER: Final = "dead-settler:1" +_CREATED_AN_HOUR_AGO: Final = ( + 'UPDATE "LiteLLM_BackgroundInteractionSettlement" ' + "SET created_at = (now() AT TIME ZONE 'UTC') - interval '1 hour' WHERE interaction_id = %s" +) _SETTLEMENT_TABLE_PRESENT_QUERY: Final = "SELECT to_regclass(%s) IS NOT NULL AS present" _SETTLEMENT_TABLE: Final = '"LiteLLM_BackgroundInteractionSettlement"' _SETTLEMENT_BY_CALL_QUERY: Final = ( @@ -179,6 +197,40 @@ def _write_config(directory: Path, upstream_url: str, models: Deployments) -> Pa return path +def _variant_config(rig: Rig, name: str, document: Mapping[str, JsonValue]) -> Path: + path: Final = rig.directory / name + path.write_text(yaml.safe_dump(dict(document))) + return path + + +def _rig_config(rig: Rig) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_python(yaml.safe_load(rig.config.read_text())) + + +def _config_without_spend_logs(rig: Rig) -> Path: + config: Final = _rig_config(rig) + general_settings: Final = {**object_value(config["general_settings"]), "disable_spend_logs": True} + return _variant_config( + rig, "settlement_config_without_spend_logs.yaml", {**config, "general_settings": general_settings} + ) + + +def _config_with_a_rotated_key(rig: Rig) -> Path: + config: Final = _rig_config(rig) + model_list: Final = [_with_rotated_key(entry, rig.models.in_progress) for entry in list_value(config["model_list"])] + return _variant_config(rig, "settlement_config_rotated_key.yaml", {**config, "model_list": model_list}) + + +def _with_rotated_key(entry: JsonValue, model_name: str) -> JsonValue: + deployment: Final = object_value(entry) + if deployment["model_name"] != model_name: + return deployment + return { + **deployment, + "litellm_params": {**object_value(deployment["litellm_params"]), "api_key": _ROTATED_DEPLOYMENT_KEY}, + } + + def _interaction(status: str, usage: dict[str, JsonValue] | None, http_status: int = 200) -> JsonResponse: return JsonResponse( content_type="application/json", @@ -206,12 +258,14 @@ def _create( path: str = "/v1beta/interactions", background: bool = True, text: str | None = None, + headers: Mapping[str, str] | None = None, ) -> str: response: Final = replica.request( "POST", path, {"model": model, "input": text or f"settle {uuid.uuid4().hex}", "background": background}, key=key, + headers=headers, ) assert response.status_code == 200, response.text return string_value(JSON_OBJECT.validate_json(response.content)["id"]) @@ -252,6 +306,16 @@ def _settlement(interaction_id: str) -> dict[str, JsonValue] | None: return rows[0] if rows else None +def _registered(interaction_id: str) -> dict[str, JsonValue]: + row: Final = eventually(lambda: _settlement(interaction_id), lambda value: value is not None) + assert row is not None + return row + + +def _abandon_claim(interaction_id: str, create_context: JsonValue) -> None: + write_rows(_ABANDONED_CLAIM, (_DEAD_SETTLER, json.dumps(create_context), interaction_id)) + + def _settlement_table_present() -> bool: return read_rows(_SETTLEMENT_TABLE_PRESENT_QUERY, (_SETTLEMENT_TABLE,))[0]["present"] is True @@ -315,6 +379,11 @@ def _assert_billed(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> return spend +def _assert_billed_with_tag(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> None: + _assert_billed(row, rates) + assert _TAG in list_value(row["request_tags"]), row + + def _key_spend(replica: Gateway, key: str) -> float: spend: Final = object_value(replica.get("/key/info", {"key": key})["info"])["spend"] assert isinstance(spend, float | int), spend @@ -333,12 +402,12 @@ def _drain(rig: Rig) -> list[JsonValue]: return requests -def _calls(rig: Rig, interaction_id: str) -> tuple[tuple[str, str], ...]: - suffix: Final = f"/v1beta/interactions/{interaction_id}" +def _calls(rig: Rig, *interaction_ids: str) -> tuple[tuple[str, str], ...]: + suffixes: Final = tuple(f"/v1beta/interactions/{interaction_id}" for interaction_id in interaction_ids) return tuple( (string_value(object_value(entry)["method"]), string_value(object_value(entry)["api_key"])) for entry in _drain(rig) - if string_value(object_value(entry)["path"]).endswith(suffix) + if string_value(object_value(entry)["path"]).endswith(suffixes) ) @@ -447,7 +516,7 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R model: Final = rig.models.in_progress key: Final = scenario.key() created_after: Final = time.time() - created: Final = tuple(_create(rig.creator, model, key) for _ in range(3)) + created: Final = tuple(_create(rig.creator, model, key, headers={"x-litellm-tags": _TAG}) for _ in range(3)) for item in created: _state(rig, item, _completed()) rates: Final = _rates(rig.creator, model) @@ -457,13 +526,113 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R pids: Final = _worker_pids(resumer.process.pid) assert len(pids) == 2, pids for item in created: - _assert_billed(_await_spend_row(item, seconds=90), rates) + _assert_billed_with_tag(_await_spend_row(item, seconds=90), rates) claimer: Final = _claimer_pid(_await_outcome(item, "billed")) assert claimer in pids or _booted_after(claimer, created_after), (claimer, pids) + assert {key for _, key in _calls(rig, *created)} == {_DEPLOYMENT_KEY} for item in created: assert len(_spend_rows(item)) == 1 +def test_a_replica_booting_later_resumes_a_database_model_with_its_own_key(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = scenario.model( + model=_MODEL, api_key=_DATABASE_DEPLOYMENT_KEY, api_base=f"{rig.upstream.url}/{rig.models.in_progress}" + ) + created: Final = _create(rig.creator, model, scenario.key()) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2 + ): + _await_spend_row(created, seconds=90) + _await_outcome(created, "billed") + assert {key for _, key in _calls(rig, created)} == {_DATABASE_DEPLOYMENT_KEY} + + +def test_a_replica_booting_later_takes_over_lapsed_claims_and_bills_only_the_unbilled(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + key: Final = scenario.key() + unbilled: Final = _create(rig.creator, model, key) + billed: Final = _create(rig.settler, model, key) + contexts: Final = {item: _registered(item)["create_context"] for item in (unbilled, billed)} + for item in (unbilled, billed): + _state(rig, item, _completed()) + rates: Final = _rates(rig.creator, model) + spend: Final = _assert_billed(_await_spend_row(billed), rates) + _await_outcome(billed, "billed") + for item, create_context in contexts.items(): + _abandon_claim(item, create_context) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2 + ): + _assert_billed(_await_spend_row(unbilled, seconds=90), rates) + assert _await_outcome(unbilled, "billed")["claimed_by"] != _DEAD_SETTLER + assert _await_outcome(billed, "billed", seconds=90)["claimed_by"] != _DEAD_SETTLER + held( + lambda: _key_spend(rig.creator, key), + lambda observed: math.isclose(observed, 2 * spend, rel_tol=1e-9), + holding=5, + seconds=30, + ) + for item in (unbilled, billed): + assert len(_spend_rows(item)) == 1 + + +def test_a_lapsed_claim_with_spend_logs_off_is_recorded_unsettled_and_never_billed(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + key: Final = scenario.key() + created: Final = _create(rig.creator, model, key) + _abandon_claim(created, _registered(created)["create_context"]) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, + rig.directory, + rig.environment(**_EAGER_RESUMER_POLL), + config=_config_without_spend_logs(rig), + workers=2, + ): + assert _await_outcome(created, "unsettled", seconds=90)["claimed_by"] != _DEAD_SETTLER + assert _spend_rows(created) == [] + assert _key_spend(rig.creator, key) == 0 + + +def test_a_resumed_interaction_whose_deployment_key_rotated_fetches_with_the_new_key(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + created: Final = _create(rig.creator, model, scenario.key()) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, + rig.directory, + rig.environment(**_EAGER_RESUMER_POLL), + config=_config_with_a_rotated_key(rig), + workers=2, + ) as resumer: + assert _model_info(resumer.gateway, model)["id"] != _model_info(rig.creator, model)["id"] + _assert_billed(_await_spend_row(created, seconds=90), _rates(rig.creator, model)) + _await_outcome(created, "billed") + assert {key for _, key in _calls(rig, created)} == {_ROTATED_DEPLOYMENT_KEY} + + +def test_a_late_restart_retries_a_failed_fetch_within_a_full_poll_window(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + created: Final = _create(rig.creator, model, scenario.key()) + _registered(created) + write_rows(_CREATED_AN_HOUR_AGO, (created,)) + _state(rig, created, InteractionState(status="completed", usage=_USAGE, get_status=503)) + _drain(rig) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_LATE_RESUMER_POLL), config=rig.config, workers=2 + ): + eventually(lambda: _calls(rig, created), lambda calls: calls != (), seconds=90) + _state(rig, created, _completed()) + _assert_billed(_await_spend_row(created), _rates(rig.creator, model)) + _await_outcome(created, "billed") + + def test_deletes_on_the_creating_proxy_bill_each_interaction_once(rig: Rig) -> None: with rig.settler.scenario() as scenario: model: Final = rig.models.in_progress diff --git a/tests/unit/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py index 9dc71a4fa1b..d55994dfb19 100644 --- a/tests/unit/interactions/test_background_cost_polling.py +++ b/tests/unit/interactions/test_background_cost_polling.py @@ -1,11 +1,14 @@ import asyncio +import json import time -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from itertools import islice from typing import Optional +import httpx import pytest +import litellm from litellm.interactions.background_cost_polling import ( _create_context, _poll_intervals, @@ -21,6 +24,7 @@ from litellm.interactions.background_cost_polling import ( ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.interactions import InteractionsAPIResponse USAGE_BLOCK = { @@ -74,6 +78,7 @@ def _context( logging_obj: LitellmLogging, timeout_seconds: float = 1.0, store: Optional[InMemoryBackgroundSettlementStore] = None, + resumed: bool = False, ) -> BackgroundInteractionPollContext: return BackgroundInteractionPollContext( interaction_id="interactions/bg-abc", @@ -82,7 +87,9 @@ def _context( initial_interval_seconds=0.001, max_interval_seconds=0.002, timeout_seconds=timeout_seconds, + outcome_write_timeout_seconds=0.01, store=store if store is not None else InMemoryBackgroundSettlementStore(), + resumed=resumed, ) @@ -255,6 +262,180 @@ async def test_poller_retries_after_fetch_error_and_still_bills(): assert logging_obj.model_call_details["response_cost"] > 0 +def _provider_error_response(status_code: int) -> httpx.Response: + return httpx.Response(status_code, request=httpx.Request("GET", "https://generativelanguage.googleapis.com")) + + +@pytest.mark.parametrize( + "error", + [ + litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.PermissionDeniedError( + message="permission denied", + model="gemini-2.5-flash", + llm_provider="gemini", + response=_provider_error_response(403), + ), + litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_poll_stops_fetching_at_the_first_error_no_retry_changes_and_releases_the_reservation(error): + """ + The regression: a resumed fetch the provider refuses for good, like an + unknown id or a key rotated away, was retried every interval until the + poll timed out. + """ + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, calls = _fetch_sequence(error) + + outcome = await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.05, resumed=True), fetch_interaction=fetch + ) + + assert outcome == "unsettled" + assert len(calls) == 1 + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_a_resumed_poll_the_provider_refuses_leaves_the_bill_to_the_replica_still_polling_it(): + """ + The regression: a replica that booted with credentials unable to read the + interaction settled it as unsettled on its first 404, and the replica that + created it, still polling, then found it settled and never billed it. + """ + store = InMemoryBackgroundSettlementStore() + creator_logging_obj = _logging_obj() + resumed_fetched = asyncio.Event() + + async def creator_fetch(context): + if resumed_fetched.is_set(): + return _response("completed", with_usage=True) + return _response("in_progress", with_usage=False) + + async def resumed_fetch(context): + resumed_fetched.set() + raise litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini") + + creator_outcome, resumed_outcome = await asyncio.gather( + poll_and_log_background_interaction_cost( + _context(creator_logging_obj, store=store), fetch_interaction=creator_fetch + ), + poll_and_log_background_interaction_cost( + _context(_logging_obj(), store=store, resumed=True), fetch_interaction=resumed_fetch + ), + ) + + assert (creator_outcome, resumed_outcome) == ("billed", None) + assert creator_logging_obj.model_call_details["response_cost"] > 0 + + +@pytest.mark.parametrize( + "error", + [ + litellm.RateLimitError(message="quota exceeded", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.ServiceUnavailableError(message="overloaded", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.InternalServerError(message="internal error", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.Timeout(message="timed out", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"), + ], +) +@pytest.mark.asyncio +async def test_the_creating_poll_retries_through_any_provider_error(error): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(error, _response("completed", with_usage=True)) + + outcome = await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert outcome == "billed" + assert len(calls) == 2 + + +class _OutcomeWritesFail: + def __init__(self): + self.store = InMemoryBackgroundSettlementStore() + + async def register(self, pending): + await self.store.register(pending) + + async def pending(self, interaction_id): + return await self.store.pending(interaction_id) + + async def is_claimed(self, interaction_id): + return await self.store.is_claimed(interaction_id) + + async def claim(self, interaction_id): + return await self.store.claim(interaction_id) + + async def record_outcome(self, interaction_id, outcome): + raise RuntimeError("database unavailable") + + async def unsettled(self): + return await self.store.unsettled() + + +@pytest.mark.parametrize( + "last_status, timeout_seconds", + [("failed", 1.0), ("in_progress", 0.01)], +) +@pytest.mark.asyncio +async def test_a_release_it_could_not_record_leaves_the_reservation_to_the_settler_that_takes_over( + last_status, timeout_seconds +): + """ + The regression: a settler released the reservation and then failed to + record it, so the replica that took over the lapsed claim released the + same reservation again and the key's budget lost that amount twice. + """ + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response(last_status, with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=timeout_seconds, store=_OutcomeWritesFail()), + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is False + + +class _FirstOutcomeWriteFails(_OutcomeWritesFail): + def __init__(self): + super().__init__() + self.failed_once = False + + async def record_outcome(self, interaction_id, outcome): + if not self.failed_once: + self.failed_once = True + raise RuntimeError("connection reset") + await self.store.record_outcome(interaction_id, outcome) + + +@pytest.mark.asyncio +async def test_a_release_whose_first_outcome_write_failed_still_frees_the_reservation(): + """ + The regression: one failed outcome write left the key's budget + reservation held until a later boot or delete took the row over, and an + active key's spend counter never expired it in the meantime. + """ + reservation = _reservation() + store = _FirstOutcomeWriteFails() + fetch, _ = _fetch_sequence(_response("failed", with_usage=False)) + + outcome = await poll_and_log_background_interaction_cost( + _context(_logging_obj_with_reservation(reservation), store=store), fetch_interaction=fetch + ) + + assert outcome == "released" + assert reservation["finalized"] is True + assert await store.store.unsettled() == () + + @pytest.mark.asyncio async def test_schedule_creates_poll_task_for_in_progress_create(): logging_obj = _logging_obj() @@ -302,6 +483,59 @@ async def test_schedule_registers_an_agent_only_create_that_names_no_model(): assert _rebuild_logging_obj(pending.create_context).model is None +@pytest.mark.asyncio +async def test_schedule_registers_a_create_whose_tags_are_a_string_and_keeps_them_as_sent(): + """ + The regression: storing the create's context read string tags as a list + and raised, so a create sent with a tags string answered with an error + after the provider had already started the interaction. + """ + logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata(tags="team-a,nightly-batch")}) + store = InMemoryBackgroundSettlementStore() + + task = await maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + store=store, + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + pending = await store.pending("interactions/bg-abc") + assert pending is not None + rebuilt = _rebuild_logging_obj(pending.create_context) + assert get_litellm_metadata_from_kwargs(kwargs=rebuilt.model_call_details)["tags"] == "team-a,nightly-batch" + + +@pytest.mark.asyncio +async def test_delete_fetches_with_the_vertex_project_location_and_credentials_the_request_used(): + store = InMemoryBackgroundSettlementStore() + await _create_on_a_replica_that_then_dies(_logging_obj(litellm_params={"metadata": _create_metadata()}), store) + fetch, captured = _capturing_fetch(_response("completed", with_usage=True)) + + outcome = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + delete_kwargs={ + "vertex_project": "request-project", + "vertex_location": "europe-west4", + "vertex_credentials": {"type": "service_account", "project_id": "request-project"}, + }, + fetch_interaction=fetch, + store=store, + ) + + assert outcome == "billed" + credentials = captured[0].credentials + assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == ( + "request-project", + "europe-west4", + {"type": "service_account", "project_id": "request-project"}, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "response,create_kwargs", @@ -781,6 +1015,105 @@ async def test_restart_resumes_only_the_rows_no_replica_claimed(): assert await store.is_claimed("interactions/bg-orphaned") +@pytest.mark.asyncio +async def test_a_restart_long_after_the_create_still_retries_a_failed_fetch_before_giving_up(): + """ + The regression: a resumed poll kept only what was left of the original + window, so a row older than the window got a single fetch, and one + transient error right after the restart ended it unbilled. + """ + store = InMemoryBackgroundSettlementStore() + await store.register( + PendingBackgroundInteraction( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"), + created_at=datetime.now(timezone.utc) - timedelta(seconds=FAST_SCHEDULE.timeout_seconds * 100), + ) + ) + fetch, calls = _fetch_sequence( + RuntimeError("transient network error"), + _response("completed", with_usage=True), + ) + + (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert len(calls) == 2 + + +def _create_through_the_proxy_router(tags: list) -> LitellmLogging: + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=False, + call_type="acreate_interaction", + start_time=time.time(), + litellm_call_id="bg-interactions-call-id", + function_id="bg-interactions-fn-id", + ) + logging_obj.update_from_kwargs( + kwargs={ + "metadata": _create_metadata( + tags=tags, spend_logs_metadata={"job": "nightly"}, model_group="client-chosen-group" + ), + "litellm_metadata": {"model_info": {"id": "deployment-1"}, "model_group": "gemini-background"}, + }, + model="gemini-2.5-flash", + optional_params={}, + litellm_params={"litellm_call_id": "bg-interactions-call-id"}, + custom_llm_provider="gemini", + ) + return logging_obj + + +def _create_with_header_tags(tags: list) -> LitellmLogging: + return _logging_obj( + litellm_params={ + "metadata": _create_metadata(tags=tags, spend_logs_metadata={"job": "nightly"}), + "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.0"}}, + } + ) + + +async def _billed_spend_row(logging_obj: LitellmLogging) -> dict: + response = _response("completed", with_usage=True) + fetch, _ = _fetch_sequence(response) + assert await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) == "billed" + settled_at = datetime.now(timezone.utc) + row = get_logging_payload( + kwargs={**logging_obj.model_call_details, "completion_start_time": settled_at}, + response_obj=response, + start_time=settled_at, + end_time=settled_at, + ) + return { + "request_tags": json.loads(row["request_tags"]), + "model_id": row["model_id"], + "model_group": row["model_group"], + "spend_logs_metadata": json.loads(row["metadata"]).get("spend_logs_metadata"), + } + + +@pytest.mark.parametrize("create", [_create_through_the_proxy_router, _create_with_header_tags]) +@pytest.mark.asyncio +async def test_a_bill_settled_off_the_creating_worker_matches_the_creating_worker_spend_row(create): + """ + The regressions: the stored create context kept only the router's metadata + on the proxy path, so a row billed after a restart or by another replica + lost the request tags that tag-based spend reports group by; then it let + the client's request metadata override the router's, so a client-sent + model group relabeled the bill another replica wrote. + """ + creator_row = await _billed_spend_row(create(["team-a", "nightly-batch"])) + stored = _create_context(create(["team-a", "nightly-batch"]), "gemini") + + rebuilt_row = await _billed_spend_row(_rebuild_logging_obj(stored)) + + assert "team-a" in creator_row["request_tags"] + assert rebuilt_row == creator_row + + class _ClaimAnswersOnlyAfterTheLastFetch: def __init__(self): self.store = InMemoryBackgroundSettlementStore() @@ -809,8 +1142,8 @@ class _ClaimAnswersOnlyAfterTheLastFetch: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): - return await self.store.unclaimed() + async def unsettled(self): + return await self.store.unsettled() @pytest.mark.asyncio @@ -844,7 +1177,7 @@ class _DownStore: async def record_outcome(self, interaction_id, outcome): raise RuntimeError("database unavailable") - async def unclaimed(self): + async def unsettled(self): raise RuntimeError("database unavailable") @@ -868,8 +1201,8 @@ class _RegistersThenRaises: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): - return await self.store.unclaimed() + async def unsettled(self): + return await self.store.unsettled() @pytest.mark.asyncio @@ -898,7 +1231,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s assert outcome == "billed" assert await store.is_claimed("interactions/bg-abc") - assert await store.unclaimed() == () + assert await store.unsettled() == () task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -907,7 +1240,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s @pytest.mark.asyncio async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch(): """ - After a restart every worker resumes the unclaimed rows, so none of them + After a restart every worker resumes the unsettled rows, so none of them is the creator whose delete may release and delete on a failed fetch. A resumed worker's delete fails like any other replica's, and its own poll still bills the interaction once it completes. @@ -965,9 +1298,9 @@ class _LandsThenGoesDown: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): + async def unsettled(self): self._answer() - return await self.store.unclaimed() + return await self.store.unsettled() def _answer(self): if self.down: @@ -992,7 +1325,7 @@ class _TableLessStore: async def record_outcome(self, interaction_id, outcome): raise RuntimeError("the settlement table does not exist") - async def unclaimed(self): + async def unsettled(self): raise RuntimeError("the settlement table does not exist") @@ -1027,7 +1360,7 @@ async def test_create_whose_registration_and_read_back_both_failed_bills_once_th assert (while_down, recovered) == (None, "billed") assert len(calls) == 2 assert await store.is_claimed("interactions/bg-abc") - assert await store.unclaimed() == () + assert await store.unsettled() == () assert await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) == () task.cancel() with pytest.raises(asyncio.CancelledError): diff --git a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py index 21384279bbb..eea7f8a81aa 100644 --- a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py +++ b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py @@ -1,22 +1,28 @@ import asyncio import time from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Optional import pytest import litellm.interactions.background_cost_polling as bg +from litellm import Router +from litellm.constants import BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS from litellm.interactions.background_cost_polling import ( _create_context, + BackgroundInteractionPollContext, configure_background_settlement_store, maybe_settle_background_interaction_before_delete, PendingBackgroundInteraction, + poll_and_log_background_interaction_cost, PollSchedule, + resume_unsettled_background_interactions, ) from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.spend_tracking.background_interaction_settlement import ( configure_background_interaction_settlement, + fetch_with_deployment_credentials, install_background_interaction_settlement, PrismaBackgroundSettlementStore, ) @@ -78,7 +84,42 @@ class _FakeSettlementTable: return len(matched) def _matching(self, where) -> list: - return [row for row in self.rows.values() if all(getattr(row, column) == value for column, value in where.items())] + return [row for row in self.rows.values() if _matches(row, where)] + + +def _matches(row: "_Row", where: dict) -> bool: + return all(_satisfies(row, column, condition) for column, condition in where.items()) + + +def _satisfies(row: "_Row", column: str, condition) -> bool: + value = getattr(row, column) + if isinstance(condition, dict): + return value is not None and value < condition["lt"] + return value == condition + + +class _FakeSpendLogs: + """The spend log table, holding only the request ids of the rows written so far.""" + + def __init__(self, request_ids: tuple[str, ...] = ()): + self.request_ids = request_ids + + async def find_unique(self, *, where): + return where["request_id"] if where["request_id"] in self.request_ids else None + + +def _store( + table, + claimed_by: str, + claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, +) -> PrismaBackgroundSettlementStore: + return PrismaBackgroundSettlementStore( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by=claimed_by, + claim_lease_seconds=claim_lease_seconds, + ) def _logging_obj(metadata: Optional[dict] = None) -> LitellmLogging: @@ -110,7 +151,12 @@ def _pending(interaction_id: str) -> PendingBackgroundInteraction: ) -def _stored_row(interaction_id: str, claimed: bool = False, create_context: Optional[object] = None) -> _Row: +def _stored_row( + interaction_id: str, + claimed_at: Optional[datetime] = None, + settled: bool = False, + create_context: Optional[object] = None, +) -> _Row: return _Row( interaction_id=interaction_id, custom_llm_provider="gemini", @@ -120,11 +166,21 @@ def _stored_row(interaction_id: str, claimed: bool = False, create_context: Opti else _create_context(_logging_obj(), "gemini").model_dump(mode="json") ), created_at=datetime.now(timezone.utc), - claimed_at=datetime.now(timezone.utc) if claimed else None, - claimed_by="replica-a:1" if claimed else None, + claimed_at=claimed_at, + claimed_by="replica-a:1" if claimed_at is not None else None, + settled_at=claimed_at if settled else None, + outcome="billed" if settled else None, ) +def _minutes_ago(minutes: float) -> datetime: + return datetime.now(timezone.utc) - timedelta(minutes=minutes) + + +def _past_the_default_lease() -> datetime: + return datetime.now(timezone.utc) - timedelta(seconds=BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS + 60) + + def _completed(interaction_id: str) -> InteractionsAPIResponse: return InteractionsAPIResponse( id=interaction_id, model="gemini-2.5-flash", status="completed", steps=[], usage=dict(USAGE_BLOCK) @@ -144,20 +200,20 @@ def _capturing_fetch(): @pytest.mark.asyncio async def test_registered_row_reads_back_as_the_same_pending_interaction(): table = _FakeSettlementTable() - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") + store = _store(table, claimed_by="replica-a:1") pending = _pending("interactions/bg-1") await store.register(pending) assert await store.pending("interactions/bg-1") == pending - assert await store.unclaimed() == (pending,) + assert await store.unsettled() == (pending,) @pytest.mark.asyncio async def test_claim_is_won_by_exactly_one_settler(): table = _FakeSettlementTable() - replica_a = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") - replica_b = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1") + replica_a = _store(table, claimed_by="replica-a:1") + replica_b = _store(table, claimed_by="replica-b:1") await replica_a.register(_pending("interactions/bg-1")) assert await replica_b.claim("interactions/bg-1") is True @@ -200,7 +256,7 @@ class _MissingSettlementTable: async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): from prisma.errors import TableNotFoundError - store = PrismaBackgroundSettlementStore(table=_MissingSettlementTable(), claimed_by="replica-a:1") + store = _store(_MissingSettlementTable(), claimed_by="replica-a:1") with pytest.raises(TableNotFoundError): await store.register(_pending("interactions/bg-1")) @@ -208,29 +264,62 @@ async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): assert await store.is_claimed("interactions/bg-1") is False assert await store.claim("interactions/bg-1") is False with pytest.raises(TableNotFoundError): - await store.unclaimed() + await store.unsettled() @pytest.mark.asyncio -async def test_unclaimed_skips_claimed_and_unreadable_rows(): +async def test_unsettled_lists_claimed_rows_with_no_outcome_and_when_their_lease_ends(): + claimed_at = _minutes_ago(1) table = _FakeSettlementTable( rows=( _stored_row("interactions/bg-orphaned"), - _stored_row("interactions/bg-settled", claimed=True), + _stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at), + _stored_row("interactions/bg-settled", claimed_at=claimed_at, settled=True), _stored_row("interactions/bg-from-the-future", create_context={"schema": "unknown"}), ) ) - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1") + store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300) - unclaimed = await store.unclaimed() + unsettled = await store.unsettled() - assert [row.interaction_id for row in unclaimed] == ["interactions/bg-orphaned"] + assert [(row.interaction_id, row.claim_expires_at) for row in unsettled] == [ + ("interactions/bg-orphaned", None), + ("interactions/bg-mid-settlement", claimed_at + timedelta(seconds=300)), + ] + + +@pytest.mark.asyncio +async def test_a_claim_whose_settler_died_is_taken_over_once_its_lease_runs_out(): + """ + The regression: a settler that died between its claim and its outcome left + the row claimed forever, so no replica, restart, or delete ever billed it. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-dead-settler", claimed_at=_minutes_ago(10)), + _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(10), settled=True), + ) + ) + replica_b = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300) + replica_c = _store(table, claimed_by="replica-c:1", claim_lease_seconds=300) + + assert await replica_b.is_claimed("interactions/bg-dead-settler") is False + assert await replica_b.pending("interactions/bg-dead-settler") is not None + assert await replica_b.claim("interactions/bg-dead-settler") is True + assert await replica_c.claim("interactions/bg-dead-settler") is False + assert await replica_c.is_claimed("interactions/bg-dead-settler") is True + assert table.rows["interactions/bg-dead-settler"].claimed_by == "replica-b:1" + + assert await replica_b.is_claimed("interactions/bg-settled") is True + assert await replica_b.pending("interactions/bg-settled") is None + assert await replica_b.claim("interactions/bg-settled") is False + assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1" @pytest.mark.asyncio async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request_context(): table = _FakeSettlementTable() - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") + store = _store(table, claimed_by="replica-a:1") await store.register(_pending("interactions/bg-1")) assert await store.claim("interactions/bg-1") assert table.rows["interactions/bg-1"].create_context @@ -244,23 +333,198 @@ async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request assert row.create_context == {} +@pytest.mark.asyncio +async def test_a_takeover_records_the_bill_a_settler_wrote_before_dying_instead_of_billing_again(): + """ + The regression: a settler that billed and then died, or failed to record + the outcome, left its claim to lapse with no outcome, and the replica that + took the claim over billed the interaction a second time. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-billed-then-died", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-died-before-billing", claimed_at=_past_the_default_lease()), + ) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + resumed = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(request_ids=("interactions/bg-billed-then-died",)), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"] + assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [ + ("replica-b:1", "billed"), + ("replica-b:1", "billed"), + ] + assert await bg._STORE.store.unsettled() == () + finally: + configure_background_settlement_store(previous_store) + + +@pytest.mark.asyncio +async def test_with_spend_logs_off_a_takeover_leaves_the_interaction_unbilled_instead_of_billing_it_twice(): + """ + The regression: with spend logs off a settler that billed and then died + leaves no spend row to find, so the replica that took over its lapsed + claim billed the interaction a second time. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-claimed-then-died", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-orphaned"), + ) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + resumed = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: False, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"] + assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [ + ("replica-b:1", "unsettled"), + ("replica-b:1", "billed"), + ] + finally: + configure_background_settlement_store(previous_store) + + +class _SpendLogsWhoseFirstReadFails(_FakeSpendLogs): + def __init__(self): + super().__init__() + self.reads = 0 + + async def find_unique(self, *, where): + self.reads += 1 + if self.reads == 1: + raise RuntimeError("read replica unavailable") + return await super().find_unique(where=where) + + +@pytest.mark.asyncio +async def test_a_takeover_whose_spend_log_read_failed_retries_and_bills_instead_of_stranding_the_row(): + """ + The regression: a takeover wrote its claim before reading the spend log, + so a failed read left the row under a fresh claim of its own, which the + poll then took for another settler's and left unbilled until a reboot. + """ + table = _FakeSettlementTable( + rows=(_stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()),) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + (resumed,) = await configure_background_interaction_settlement( + table=table, + spend_logs=_SpendLogsWhoseFirstReadFails(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert table.rows["interactions/bg-dead-settler"].outcome == "billed" + finally: + configure_background_settlement_store(previous_store) + + +def _logging_obj_whose_bill_takes(seconds: float, bills: list) -> LitellmLogging: + logging_obj = _logging_obj() + + async def bill(result): + bills.append(result.id) + await asyncio.sleep(seconds) + + logging_obj.async_log_background_interaction_completion = bill + return logging_obj + + +@pytest.mark.asyncio +async def test_a_settle_that_outlasts_the_lease_is_not_taken_over_and_billed_twice(): + """ + The regression: a success callback that held the settle past the claim + lease let a booting replica take the live claim over and bill the same + interaction a second time. + """ + lease_seconds = 0.6 + table = _FakeSettlementTable() + creator = _store(table, claimed_by="replica-a:1", claim_lease_seconds=lease_seconds) + resumer = _store(table, claimed_by="replica-b:1", claim_lease_seconds=lease_seconds) + await creator.register(_pending("interactions/bg-slow-settle")) + fetch, _ = _capturing_fetch() + bills = [] + + def poll(store, resumed: bool, start_delay_seconds: float): + return poll_and_log_background_interaction_cost( + BackgroundInteractionPollContext( + interaction_id="interactions/bg-slow-settle", + custom_llm_provider="gemini", + logging_obj=_logging_obj_whose_bill_takes(lease_seconds * 2.5, bills), + initial_interval_seconds=0.001, + max_interval_seconds=0.002, + timeout_seconds=lease_seconds * 3, + start_delay_seconds=start_delay_seconds, + claim_renewal_interval_seconds=lease_seconds / 10, + store=store, + resumed=resumed, + ), + fetch, + ) + + outcomes = await asyncio.gather( + poll(creator, resumed=False, start_delay_seconds=0.0), + poll(resumer, resumed=True, start_delay_seconds=lease_seconds * 1.5), + ) + + assert outcomes == ["billed", None] + assert bills == ["interactions/bg-slow-settle"] + assert table.rows["interactions/bg-slow-settle"].claimed_by == "replica-a:1" + + @pytest.mark.asyncio async def test_configure_installs_the_store_and_resumes_the_orphaned_rows(): table = _FakeSettlementTable( - rows=(_stored_row("interactions/bg-orphaned"), _stored_row("interactions/bg-settled", claimed=True)) + rows=( + _stored_row("interactions/bg-orphaned"), + _stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(1), settled=True), + ) ) fetch, captured = _capturing_fetch() previous_store = bg._STORE.store try: resumed = await configure_background_interaction_settlement( - table=table, claimed_by="replica-b:1", fetch_interaction=fetch, schedule=FAST_SCHEDULE + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, ) - assert len(resumed) == 1 - assert await asyncio.wait_for(resumed[0], timeout=5) == "billed" - assert [context.interaction_id for context in captured] == ["interactions/bg-orphaned"] - assert table.rows["interactions/bg-orphaned"].claimed_by == "replica-b:1" - assert table.rows["interactions/bg-orphaned"].outcome == "billed" + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == ["billed", "billed"] + assert sorted(context.interaction_id for context in captured) == [ + "interactions/bg-dead-settler", + "interactions/bg-orphaned", + ] + for interaction_id in ("interactions/bg-orphaned", "interactions/bg-dead-settler"): + assert table.rows[interaction_id].claimed_by == "replica-b:1" + assert table.rows[interaction_id].outcome == "billed" + assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1" await table.create( data={ @@ -288,7 +552,9 @@ class _PrismaClientWithoutSettlementTable: async def test_install_keeps_booting_when_the_settlement_table_is_unreachable(): previous_store = bg._STORE.store - await install_background_interaction_settlement(_PrismaClientWithoutSettlementTable()) + await install_background_interaction_settlement( + _PrismaClientWithoutSettlementTable(), router=lambda: None, spend_logs_kept=lambda: True + ) assert bg._STORE.store is previous_store @@ -296,3 +562,166 @@ async def test_install_keeps_booting_when_the_settlement_table_is_unreachable(): @dataclass(frozen=True) class _JsonLike: data: object + + +@pytest.mark.asyncio +async def test_a_resumed_row_still_under_a_live_claim_waits_for_the_lease_before_settling(): + claimed_at = datetime.now(timezone.utc) + table = _FakeSettlementTable(rows=(_stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at),)) + store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=0.2) + fetched_at = [] + + async def fetch(context): + fetched_at.append(datetime.now(timezone.utc)) + return _completed(context.interaction_id) + + (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert fetched_at[0] >= claimed_at + timedelta(seconds=0.2) + assert table.rows["interactions/bg-mid-settlement"].claimed_by == "replica-b:1" + + +def _router_with_a_deployment_scoped_key() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-background", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "deployment-scoped-key", + "api_base": "https://deployment-scoped.example", + }, + "model_info": {"id": "deployment-1"}, + } + ] + ) + + +def _router_with_a_catch_all_route() -> Router: + return Router(model_list=[{"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "catch-all-key"}}]) + + +def _router_with_two_deployments_in_the_group() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-background", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": f"deployment-key-{index}"}, + "model_info": {"id": f"deployment-{index}"}, + } + for index in (2, 3) + ] + ) + + +def _router_with_a_vertex_deployment() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-background", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "deployment-project", + "vertex_location": "europe-west4", + "vertex_credentials": "/secrets/deployment-service-account.json", + }, + "model_info": {"id": "vertex-deployment-1"}, + } + ] + ) + + +def _create_context_served_by(deployment_id: str) -> dict: + metadata = { + "user_api_key": "0123456789abcdef" * 4, + "model_info": {"id": deployment_id}, + "model_group": "gemini-background", + } + return _create_context(_logging_obj(metadata), "gemini").model_dump(mode="json") + + +@pytest.mark.parametrize("deployment_id", ["deployment-1", "deployment-1-before-an-edit"]) +@pytest.mark.asyncio +async def test_a_resumed_poll_fetches_with_the_credentials_of_the_deployment_that_served_the_create(deployment_id): + """ + The regression: a resumed poll has no request to take credentials from, so + it fetched with the environment's key, which a deployment-scoped key is + not, and every fetch failed until the row was given up unbilled. Editing + the deployment since gives it a new id, and its model group still names it. + """ + table = _FakeSettlementTable( + rows=(_stored_row("interactions/bg-1", create_context=_create_context_served_by(deployment_id)),) + ) + router = _router_with_a_deployment_scoped_key() + fetch, captured = _capturing_fetch() + previous_store = bg._STORE.store + try: + (resumed,) = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch_with_deployment_credentials(lambda: router, fetch), + schedule=FAST_SCHEDULE, + ) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert (captured[0].credentials.api_key, captured[0].credentials.api_base) == ( + "deployment-scoped-key", + "https://deployment-scoped.example", + ) + finally: + configure_background_settlement_store(previous_store) + + +@pytest.mark.parametrize( + "router, deployment_id, model_group", + [ + (_router_with_a_deployment_scoped_key, "deployment-removed-since", None), + (_router_with_a_deployment_scoped_key, None, None), + (lambda: None, "deployment-1", "gemini-background"), + (_router_with_a_catch_all_route, "deployment-removed-since", None), + (_router_with_two_deployments_in_the_group, "deployment-removed-since", "gemini-background"), + (_router_with_a_vertex_deployment, "deployment-removed-since", "vertex-background"), + (_router_with_a_vertex_deployment, "vertex-deployment-1", None), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_poll_without_a_known_deployment_fetches_with_the_environment_credentials( + router, deployment_id, model_group +): + live_router = router() + fetch, captured = _capturing_fetch() + context = BackgroundInteractionPollContext( + interaction_id="interactions/bg-1", + custom_llm_provider="gemini", + logging_obj=_logging_obj(), + deployment_id=deployment_id, + model_group=model_group, + ) + + await fetch_with_deployment_credentials(lambda: live_router, fetch)(context) + + assert captured == [context] + + +@pytest.mark.asyncio +async def test_a_resumed_vertex_poll_fetches_with_the_project_location_and_credentials_of_its_deployment(): + router = _router_with_a_vertex_deployment() + fetch, captured = _capturing_fetch() + context = BackgroundInteractionPollContext( + interaction_id="interactions/bg-1", + custom_llm_provider="vertex_ai", + logging_obj=_logging_obj(), + deployment_id="vertex-deployment-1", + ) + + await fetch_with_deployment_credentials(lambda: router, fetch)(context) + + credentials = captured[0].credentials + assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == ( + "deployment-project", + "europe-west4", + "/secrets/deployment-service-account.json", + )