mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(interactions): keep background interaction bills through crashes, late restarts, and deployment keys (#45743)
* fix(interactions): keep background interaction bills through crashes, late restarts, and deployment keys A settlement claim is now a lease, so a row whose settler died between claim and outcome is claimable again after the lease and boot resumes it. A resumed poll gets the full poll window and fetches with the serving deployment's credentials from the live router. The stored create context keeps the request tags and the request metadata that the creating worker's own spend row used * fix(interactions): bill a lapsed claim's takeover only when no spend row exists A settler that billed and then died, or failed to record the outcome, left the row unsettled, so the replica that took over its lapsed claim billed the interaction a second time. A takeover now looks up the spend row keyed by the interaction id and records the row as billed when it finds one. * fix(interactions): settle Vertex AI and edited deployments, never double bill with spend logs off A takeover with spend logs off records the interaction unsettled instead of billing it again, and a release records its outcome before freeing the budget reservation. A resumed or live fetch carries the deployment's Vertex AI project, location, and credentials, and a stale deployment id falls back to the model group only when one same-provider deployment serves it. A poll stops at the first 400, 401, 403, or 404, string request tags no longer fail the create, and settled_at gets an index built concurrently. * fix(interactions): leave a live replica's interaction to it when a resumed poll cannot read it A replica that boots with credentials unable to read an interaction got a 404 on its resumed poll and settled the row as unsettled at once, so the replica that created it, still polling, found it settled and never billed it. A resumed poll that gets a 400, 401, 403, or 404 now stops fetching and keeps checking the row until its poll window ends, and settles it as unsettled only when no other replica did. * fix(interactions): keep the creator's metadata on a resumed bill and retry outcome writes inside the claim lease A bill another replica writes now carries the metadata the creating worker's own spend row is built from, so a client's metadata.model_group no longer relabels it. A failed outcome write is retried for half the claim lease before the budget reservation is left to a takeover, and the creating replica retries every fetch error for its whole window again, as main does. The resume integration case now checks that a resumed fetch uses the deployment's own key. * fix(interactions): resume background interaction polls only after database models load Boot resume ran before the proxy loaded the models stored in the database, so a resumed poll for a model added through /model/new or the Admin UI fetched with env credentials. A 4xx there stops a resumed poll, so the bill was lost. The settlement now installs after the database models load. The integration test resumes a database model with an eager first fetch and fails with the old startup order. * fix(interactions): renew a settling claim so a slow settle is never taken over and billed twice A settle runs every success callback in turn, and one slow callback could hold it past the 5 minute claim lease. A booting replica or a DELETE then took the live claim over, found no spend row yet, and billed the interaction a second time. The settler now renews its claim every third of the lease until it records the outcome, so only a settler that stopped renewing is taken over. A takeover also read the spend log after writing its claim, so a failed read left the row under the taker's own fresh claim, which its poll then took for another settler's and left unbilled until a reboot. The spend row is now read before the claim is written, so a failed read takes nothing and the claim is retried on the next poll. * test(spend): cover lapsed claim takeovers, rotated deployment keys, late restarts, and resumed bill tags --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
2dc6aafa6b
commit
4e2dd65549
11 changed files with 1400 additions and 124 deletions
|
|
@ -0,0 +1,2 @@
|
|||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at"
|
||||
ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at");
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue