mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
refactor(spend): fold auto-router turns in the flusher, not on the request path
Classifying a turn is a read-modify-write over the session's stored state, and the read and the modify were happening on the request path while the write happened in the background flusher. Splitting one unit of work across two contexts is what produced all three findings from the last review round. The logging path now only stages the turn's facts; record_turn no longer takes a prisma client at all, so a database round trip in front of spend processing is not expressible. The flusher owns the whole read, fold and write per session, so a read that faults has no half-finished write to corrupt. It re-stages the turns through the same path a failed write already used, and the except branch that returned an empty state, then persisted it over real history, is gone rather than guarded. Folding at flush time also lets an interval's turns be sorted by start time before they are classified. Completion order is not start order, so two turns in flight together used to leave the earlier one read as a late arrival and dropped from every bucket; at arrival there is no later turn to compare against, so this was not fixable in the previous shape. StateUnavailable separates "the read failed" from "this session has no history", which were the same value before. Only the second one is writable. The staging ceiling counts turns rather than sessions, which is the quantity that actually bounds the memory a caller sending a fresh session id per request can make the proxy hold. Benchmark window dates are typed as dates on the route, so a malformed one is rejected by the framework instead of raising inside the aggregate, and an inverted range is a 400 instead of an empty dashboard that reads as no traffic
This commit is contained in:
Tin Chi Lo
parent
c5d56e4e86
commit
8d300d98d8
9 changed files with 387 additions and 191 deletions
|
|
@ -346,7 +346,6 @@ class DBSpendUpdateWriter:
|
|||
cache_creation_tokens=_extract_cache_creation_tokens(usage_obj),
|
||||
usage_object=usage_obj,
|
||||
),
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
async def _enqueue_tool_usage_transaction(
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import time
|
|||
import traceback
|
||||
import warnings
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from types import UnionType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -16483,8 +16483,8 @@ async def get_adaptive_router_state(
|
|||
dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI takes a list of dependencies
|
||||
)
|
||||
async def get_auto_router_benchmarks(
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
):
|
||||
"""Session-level benchmarks for every configured auto-router.
|
||||
|
|
@ -16494,10 +16494,11 @@ async def get_auto_router_benchmarks(
|
|||
counterfactual baseline, and how the provider prompt cache behaved.
|
||||
|
||||
Reads the per-session rollup, never the per-request spend logs.
|
||||
``start_date`` / ``end_date`` are ``YYYY-MM-DD``; the window is clamped to
|
||||
the most recent ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes
|
||||
the window actually served. Sessions are counted whole when they were active
|
||||
in the window. Returns 404 when no auto-router is configured.
|
||||
``start_date`` / ``end_date`` are ``YYYY-MM-DD``, rejected by the framework
|
||||
when malformed; the window is clamped to the most recent
|
||||
``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes the window
|
||||
actually served. Sessions are counted whole when they were active in the
|
||||
window. Returns 404 when no auto-router is configured.
|
||||
"""
|
||||
from litellm.proxy.spend_tracking.auto_router_benchmarks import compute_benchmarks
|
||||
from litellm.proxy.spend_tracking.auto_router_sessions import auto_router_group_kinds
|
||||
|
|
@ -16509,6 +16510,13 @@ async def get_auto_router_benchmarks(
|
|||
"error": CommonProxyErrors.not_allowed_access.value
|
||||
},
|
||||
)
|
||||
if end_date < start_date:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={ # mutable-ok: HTTPException takes a dict detail
|
||||
"error": "end_date must not be earlier than start_date."
|
||||
},
|
||||
)
|
||||
if llm_router is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ One aggregate query covers every auto-router, rather than four per router.
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, timedelta
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, NamedTuple
|
||||
|
||||
|
|
@ -158,18 +158,16 @@ GROUP BY model_group
|
|||
"""
|
||||
|
||||
|
||||
def clamp_window(start_date: str, end_date: str) -> _Window:
|
||||
"""Parse the range and enforce ``start >= end - BENCHMARKS_MAX_WINDOW_DAYS``.
|
||||
def clamp_window(start_date: date, end_date: date) -> _Window:
|
||||
"""Enforce ``start >= end - BENCHMARKS_MAX_WINDOW_DAYS``.
|
||||
|
||||
The returned start reflects the window actually served, which the response
|
||||
echoes so the dashboard can label what it is showing rather than what it
|
||||
asked for.
|
||||
asked for. The dates arrive already parsed, because a malformed one is the
|
||||
route's contract to reject rather than this module's to discover.
|
||||
"""
|
||||
start = datetime.fromisoformat(start_date).replace(tzinfo=timezone.utc)
|
||||
end = datetime.fromisoformat(end_date).replace(tzinfo=timezone.utc)
|
||||
floor = (end - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)).replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
clamped = max(start, floor)
|
||||
return _Window(start=clamped.date().isoformat(), end=end.date().isoformat())
|
||||
floor = end_date - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)
|
||||
return _Window(start=max(start_date, floor).isoformat(), end=end_date.isoformat())
|
||||
|
||||
|
||||
def _rate_pct(part: int, whole: int) -> float:
|
||||
|
|
@ -259,8 +257,8 @@ def summarize_group(row: _GroupRow, router_kind: str) -> AutoRouterGroupBenchmar
|
|||
async def compute_benchmarks(
|
||||
prisma_client: "PrismaClient",
|
||||
group_kinds: Mapping[str, str],
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
start_date: date,
|
||||
end_date: date,
|
||||
) -> AutoRouterBenchmarksResponse:
|
||||
"""Aggregate the session rollup for every configured auto-router."""
|
||||
window = clamp_window(start_date, end_date)
|
||||
|
|
|
|||
|
|
@ -1,18 +1,19 @@
|
|||
"""In-memory aggregation and durable flush for auto-router session rollups.
|
||||
"""In-memory staging and durable flush for auto-router session rollups.
|
||||
|
||||
Follows ``AdaptiveRouterUpdateQueue``: the logging path only folds into memory,
|
||||
and a background task drains the aggregate into Postgres with atomic increment
|
||||
upserts, so two pods writing the same session compose rather than overwrite.
|
||||
Follows ``AdaptiveRouterUpdateQueue``: the logging path only stages into memory
|
||||
and never touches the database, and a background task drains what it staged into
|
||||
Postgres with atomic increment upserts, so two pods writing the same session
|
||||
compose rather than overwrite.
|
||||
|
||||
The one departure is that this queue also caches the session state the fold reads
|
||||
from. A pod that has never seen a session loads its row once and classifies from
|
||||
memory thereafter, which is what keeps a session correct across a restart or a
|
||||
move between pods without paying a read per turn.
|
||||
Classifying a turn depends on the session's prior state, so the fold happens in
|
||||
the flusher rather than at arrival. That is the one place the stored state can be
|
||||
read, folded onto and written back as a single unit of work; a read that fails
|
||||
there has no half-finished write to corrupt, and the turns simply stage again.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, replace
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING
|
||||
|
|
@ -21,11 +22,12 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.proxy.spend_tracking.auto_router_sessions import (
|
||||
EMPTY_SESSION_STATE,
|
||||
SessionState,
|
||||
StateLookup,
|
||||
StateUnavailable,
|
||||
TurnDelta,
|
||||
TurnFacts,
|
||||
counters_of,
|
||||
fold_turn,
|
||||
merge_deltas,
|
||||
fold_session,
|
||||
state_column,
|
||||
state_from_row,
|
||||
)
|
||||
|
|
@ -36,50 +38,39 @@ if TYPE_CHECKING:
|
|||
|
||||
SessionKey = tuple[str, str]
|
||||
|
||||
DEFAULT_MAX_TRACKED_SESSIONS = 10_000
|
||||
DEFAULT_MAX_STAGED_TURNS = 50_000
|
||||
MAX_CACHED_SESSIONS = 10_000
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Pending:
|
||||
router_kind: str
|
||||
baseline_model: str | None
|
||||
first_turn_at: float
|
||||
last_turn_at: float
|
||||
delta: TurnDelta
|
||||
turns: tuple[TurnFacts, ...]
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _warn_pending_full(cap: int) -> None:
|
||||
def _warn_staging_full(cap: int) -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
"auto_router_sessions: %d sessions staged for the next flush; new sessions are not being recorded "
|
||||
"auto_router_sessions: %d turns staged for the next flush; further turns are not being recorded "
|
||||
"until it drains. Benchmarks will undercount until then",
|
||||
cap,
|
||||
)
|
||||
|
||||
|
||||
def _merge_pending(earlier: _Pending, later: _Pending) -> _Pending:
|
||||
"""Fold two staged batches for one session, oldest first."""
|
||||
return _Pending(
|
||||
router_kind=later.router_kind,
|
||||
baseline_model=later.baseline_model or earlier.baseline_model,
|
||||
first_turn_at=min(earlier.first_turn_at, later.first_turn_at),
|
||||
last_turn_at=max(earlier.last_turn_at, later.last_turn_at),
|
||||
delta=merge_deltas(earlier.delta, later.delta),
|
||||
)
|
||||
|
||||
|
||||
def _epoch_to_datetime(value: float) -> datetime:
|
||||
return datetime.fromtimestamp(value, tz=timezone.utc)
|
||||
|
||||
|
||||
class AutoRouterSessionQueue:
|
||||
"""Folds auto-routed turns in memory and flushes them to the session rollup."""
|
||||
"""Stages auto-routed turns in memory and folds them into the rollup on flush."""
|
||||
|
||||
def __init__(self, max_tracked_sessions: int = DEFAULT_MAX_TRACKED_SESSIONS) -> None:
|
||||
def __init__(self, max_staged_turns: int = DEFAULT_MAX_STAGED_TURNS) -> None:
|
||||
self._pending: dict[SessionKey, _Pending] = {} # mutable-ok: drained and replaced wholesale on flush
|
||||
self._state: OrderedDict[SessionKey, SessionState] = OrderedDict() # mutable-ok: bounded LRU cache
|
||||
self._staged_turns = 0
|
||||
self._lock = asyncio.Lock()
|
||||
self._max_tracked_sessions = max_tracked_sessions
|
||||
self._max_staged_turns = max_staged_turns
|
||||
|
||||
async def record_turn(
|
||||
self,
|
||||
|
|
@ -87,58 +78,95 @@ class AutoRouterSessionQueue:
|
|||
router_kind: str,
|
||||
baseline_model: str | None,
|
||||
turn: TurnFacts,
|
||||
prisma_client: "PrismaClient",
|
||||
) -> None:
|
||||
"""Classify one turn against its session and stage the increments.
|
||||
"""Stage one turn against its session. Does no I/O; the fold happens on flush.
|
||||
|
||||
The session id is caller-controlled, so the staged aggregate is capped:
|
||||
past the cap a session that is already staged keeps accumulating, but a
|
||||
new one is dropped rather than admitted. Without that bound a caller
|
||||
sending a fresh id per request grows the aggregate without limit between
|
||||
flushes, and benchmark rows are not worth an out-of-memory kill.
|
||||
``session_id`` is caller-controlled, so the staging is capped on turns
|
||||
held rather than on sessions seen, which is the quantity that actually
|
||||
bounds the memory. Past the cap a turn is dropped and logged, because
|
||||
benchmark rows are not worth an out-of-memory kill.
|
||||
"""
|
||||
loaded = await self._session_state(key, prisma_client)
|
||||
async with self._lock:
|
||||
current = self._pending.get(key)
|
||||
if current is None and len(self._pending) >= self._max_tracked_sessions:
|
||||
_warn_pending_full(self._max_tracked_sessions)
|
||||
if self._staged_turns >= self._max_staged_turns:
|
||||
_warn_staging_full(self._max_staged_turns)
|
||||
return
|
||||
cached = self._state.get(key)
|
||||
delta = fold_turn(cached if cached is not None else loaded, turn)
|
||||
self._remember(key, delta.state)
|
||||
self._pending[key] = (
|
||||
_Pending(
|
||||
router_kind=router_kind,
|
||||
baseline_model=baseline_model,
|
||||
first_turn_at=turn.started_at,
|
||||
last_turn_at=turn.started_at,
|
||||
delta=delta,
|
||||
)
|
||||
if current is None
|
||||
else replace(
|
||||
current,
|
||||
baseline_model=baseline_model or current.baseline_model,
|
||||
first_turn_at=min(current.first_turn_at, turn.started_at),
|
||||
last_turn_at=max(current.last_turn_at, turn.started_at),
|
||||
delta=merge_deltas(current.delta, delta),
|
||||
)
|
||||
)
|
||||
self._stage(key, router_kind, baseline_model, (turn,))
|
||||
|
||||
async def _session_state(self, key: SessionKey, prisma_client: "PrismaClient") -> SessionState:
|
||||
"""The session's state, from memory when this pod has seen it before.
|
||||
def _stage(
|
||||
self, key: SessionKey, router_kind: str, baseline_model: str | None, turns: tuple[TurnFacts, ...]
|
||||
) -> None:
|
||||
"""Add turns to whatever this session already has staged; callers hold the lock.
|
||||
|
||||
Loading outside the lock keeps a slow read from stalling every other
|
||||
session's fold; a concurrent loader for the same key at worst repeats the
|
||||
read, since both resolve to the same stored row.
|
||||
Arriving turns and a replayed batch stage identically, because the fold
|
||||
sorts by start time and so does not care which of the two came first.
|
||||
"""
|
||||
current = self._pending.get(key)
|
||||
self._pending[key] = _Pending(
|
||||
router_kind=router_kind,
|
||||
baseline_model=baseline_model or (current.baseline_model if current is not None else None),
|
||||
turns=(current.turns if current is not None else ()) + turns,
|
||||
)
|
||||
self._staged_turns += len(turns)
|
||||
|
||||
async def flush(self, prisma_client: "PrismaClient") -> int:
|
||||
"""Fold and write every staged session. Returns rows written.
|
||||
|
||||
A session that could not be read or written is staged again rather than
|
||||
dropped. Reading, folding and writing are one unit here, so a failure
|
||||
means nothing landed and replaying it cannot double-count; draining first
|
||||
and swallowing the error would lose that interval's turns, tokens and
|
||||
spend permanently on any transient database fault.
|
||||
"""
|
||||
async with self._lock:
|
||||
cached = self._state.get(key)
|
||||
if cached is not None:
|
||||
self._state.move_to_end(key)
|
||||
return cached
|
||||
return await self._load_state(key, prisma_client)
|
||||
batch = self._pending
|
||||
self._pending = {} # mutable-ok: fresh staging for the next interval
|
||||
self._staged_turns = 0
|
||||
|
||||
async def _load_state(self, key: SessionKey, prisma_client: "PrismaClient") -> SessionState:
|
||||
failed = { # mutable-ok: built once from the sessions that did not land
|
||||
key: batch[key] for key in sorted(batch.keys()) if not await self._commit(key, batch[key], prisma_client)
|
||||
}
|
||||
if failed:
|
||||
verbose_proxy_logger.warning(
|
||||
"auto_router_sessions: %d of %d sessions could not be folded; re-staging them for the next flush",
|
||||
len(failed),
|
||||
len(batch),
|
||||
)
|
||||
async with self._lock:
|
||||
for key, pending in failed.items():
|
||||
self._stage(key, pending.router_kind, pending.baseline_model, pending.turns)
|
||||
return len(batch) - len(failed)
|
||||
|
||||
async def _commit(self, key: SessionKey, pending: _Pending, prisma_client: "PrismaClient") -> bool:
|
||||
"""Read the session's state, fold its staged turns onto it, write both back.
|
||||
|
||||
The cache advances only once the write has landed, so a replay folds from
|
||||
the state the failed attempt did rather than one that was never persisted.
|
||||
An evicted session is not lost either; its next flush reloads the row it
|
||||
was already written to, which costs one read and folds identically.
|
||||
"""
|
||||
state = await self._session_state(key, prisma_client)
|
||||
if isinstance(state, StateUnavailable):
|
||||
return False
|
||||
delta = fold_session(state, pending.turns)
|
||||
if not await self._write(key, pending, delta, prisma_client):
|
||||
return False
|
||||
self._state[key] = delta.state
|
||||
self._state.move_to_end(key)
|
||||
while len(self._state) > MAX_CACHED_SESSIONS:
|
||||
self._state.popitem(last=False)
|
||||
return True
|
||||
|
||||
async def _session_state(self, key: SessionKey, prisma_client: "PrismaClient") -> StateLookup:
|
||||
"""The session's state: from memory if this pod has folded it before, else its row.
|
||||
|
||||
Only the flusher touches this cache, so it needs no lock. A session this
|
||||
pod has not seen costs one read, which is what keeps it correct across a
|
||||
restart or a move between pods.
|
||||
"""
|
||||
cached = self._state.get(key)
|
||||
if cached is not None:
|
||||
self._state.move_to_end(key)
|
||||
return cached
|
||||
session_id, model_group = key
|
||||
try:
|
||||
row = await AutoRouterSessionRepository(prisma_client).table.find_unique(
|
||||
|
|
@ -149,63 +177,20 @@ class AutoRouterSessionQueue:
|
|||
}
|
||||
}
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a read fault must not fail the spend write
|
||||
verbose_proxy_logger.warning(
|
||||
"auto_router_sessions: could not load session state for %s (%s); treating as a new session", key, e
|
||||
)
|
||||
return EMPTY_SESSION_STATE
|
||||
except Exception as e: # noqa: BLE001 # a read fault re-stages the turns rather than failing the flush
|
||||
verbose_proxy_logger.warning("auto_router_sessions: could not load session state for %s (%s)", key, e)
|
||||
return StateUnavailable()
|
||||
if row is None:
|
||||
return EMPTY_SESSION_STATE
|
||||
return state_from_row(row.last_model, row.last_turn_at, row.model_state)
|
||||
|
||||
def _remember(self, key: SessionKey, state: SessionState) -> None:
|
||||
"""Cache the session's next state, evicting the least recently used.
|
||||
|
||||
An evicted session is not lost; its next turn reloads the row it was
|
||||
already flushed to, which costs one read and classifies identically.
|
||||
"""
|
||||
self._state[key] = state
|
||||
self._state.move_to_end(key)
|
||||
while len(self._state) > self._max_tracked_sessions:
|
||||
self._state.popitem(last=False)
|
||||
|
||||
async def flush(self, prisma_client: "PrismaClient") -> int:
|
||||
"""Drain the aggregate into the session rollup. Returns rows written.
|
||||
|
||||
A session whose write fails is staged again rather than dropped. Draining
|
||||
first and swallowing the error would lose that interval's turns, tokens
|
||||
and spend permanently on any transient database fault, and because the
|
||||
upsert is atomic a failure means nothing landed, so replaying it cannot
|
||||
double-count.
|
||||
"""
|
||||
async with self._lock:
|
||||
batch = self._pending
|
||||
self._pending = {} # mutable-ok: fresh aggregate for the next interval
|
||||
|
||||
failed = { # mutable-ok: built once from the writes that did not land
|
||||
key: batch[key] for key in sorted(batch.keys()) if not await self._write(key, batch[key], prisma_client)
|
||||
}
|
||||
if failed:
|
||||
verbose_proxy_logger.warning(
|
||||
"auto_router_sessions: %d of %d session writes failed; re-staging them for the next flush",
|
||||
len(failed),
|
||||
len(batch),
|
||||
)
|
||||
async with self._lock:
|
||||
for key, pending in failed.items():
|
||||
current = self._pending.get(key)
|
||||
# The retried batch is older than anything staged since, so it
|
||||
# merges underneath it and the state of the newer one wins.
|
||||
self._pending[key] = pending if current is None else _merge_pending(pending, current)
|
||||
return len(batch) - len(failed)
|
||||
|
||||
async def _write(self, key: SessionKey, pending: _Pending, prisma_client: "PrismaClient") -> bool:
|
||||
async def _write(self, key: SessionKey, pending: _Pending, delta: TurnDelta, prisma_client: "PrismaClient") -> bool:
|
||||
session_id, model_group = key
|
||||
counters = counters_of(pending.delta)
|
||||
counters = counters_of(delta)
|
||||
shared = { # mutable-ok: prisma's write API takes dict payloads
|
||||
"last_turn_at": _epoch_to_datetime(pending.last_turn_at),
|
||||
"last_model": pending.delta.state.last_model,
|
||||
"model_state": state_column(pending.delta.state),
|
||||
"last_turn_at": _epoch_to_datetime(max(turn.started_at for turn in pending.turns)),
|
||||
"last_model": delta.state.last_model,
|
||||
"model_state": state_column(delta.state),
|
||||
"baseline_model": pending.baseline_model,
|
||||
}
|
||||
try:
|
||||
|
|
@ -223,7 +208,7 @@ class AutoRouterSessionQueue:
|
|||
"session_id": session_id,
|
||||
"model_group": model_group,
|
||||
"router_kind": pending.router_kind,
|
||||
"first_turn_at": _epoch_to_datetime(pending.first_turn_at),
|
||||
"first_turn_at": _epoch_to_datetime(min(turn.started_at for turn in pending.turns)),
|
||||
**shared,
|
||||
**counters,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ models. This is the same reason ``savings.py`` prices in the spend writer.
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, fields
|
||||
from datetime import datetime, timezone
|
||||
from functools import reduce
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Literal, Protocol
|
||||
|
||||
|
|
@ -80,6 +81,19 @@ EMPTY_SESSION_STATE = SessionState(
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StateUnavailable:
|
||||
"""The session's stored state could not be read, so there is nothing to fold onto.
|
||||
|
||||
Distinct from a session with no history: an absent row means start from empty
|
||||
and write the result, an unreadable one means try again rather than replace a
|
||||
history that really happened with one derived from a single turn.
|
||||
"""
|
||||
|
||||
|
||||
StateLookup = SessionState | StateUnavailable
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TurnFacts:
|
||||
"""One auto-routed request, as the spend writer sees it.
|
||||
|
|
@ -448,3 +462,17 @@ def merge_deltas(left: TurnDelta, right: TurnDelta) -> TurnDelta:
|
|||
name: getattr(left, name) + getattr(right, name) for name in COUNTER_FIELDS
|
||||
}, # mutable-ok: a JSON object is a dict by definition
|
||||
)
|
||||
|
||||
|
||||
def fold_session(state: SessionState, turns: tuple[TurnFacts, ...], rates: RateLookup = cache_rates) -> TurnDelta:
|
||||
"""Fold an interval of one session's turns onto its prior state, oldest first.
|
||||
|
||||
Ordering happens here rather than at arrival because completion order is not
|
||||
start order: turns in flight together finish in whichever order the providers
|
||||
answer, and the earlier one would otherwise read as a late arrival.
|
||||
"""
|
||||
return reduce(
|
||||
lambda folded, turn: merge_deltas(folded, fold_turn(folded.state, turn, rates)),
|
||||
sorted(turns, key=lambda turn: turn.started_at),
|
||||
TurnDelta(state=state),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import date, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -67,23 +67,25 @@ def group_row(**overrides):
|
|||
return row
|
||||
|
||||
|
||||
START = date(2026, 7, 2)
|
||||
END = date(2026, 8, 1)
|
||||
|
||||
|
||||
async def benchmarks_for(**overrides):
|
||||
prisma = _FakePrisma([group_row(**overrides)])
|
||||
return await compute_benchmarks(prisma, GROUP_KINDS, "2026-07-02", "2026-08-01")
|
||||
return await compute_benchmarks(prisma, GROUP_KINDS, START, END)
|
||||
|
||||
|
||||
class TestWindowClamping:
|
||||
def test_a_wider_request_is_clamped_to_the_maximum_window(self):
|
||||
window = clamp_window("2020-01-01", "2026-08-01")
|
||||
expected = (datetime(2026, 8, 1, tzinfo=timezone.utc) - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)).date()
|
||||
assert window.start == expected.isoformat()
|
||||
window = clamp_window(date(2020, 1, 1), END)
|
||||
assert window.start == (END - timedelta(days=BENCHMARKS_MAX_WINDOW_DAYS)).isoformat()
|
||||
|
||||
def test_a_narrower_request_is_served_as_asked(self):
|
||||
assert clamp_window("2026-07-25", "2026-08-01").start == "2026-07-25"
|
||||
assert clamp_window(date(2026, 7, 25), END).start == "2026-07-25"
|
||||
|
||||
def test_the_response_echoes_the_window_actually_served(self):
|
||||
window = clamp_window("2020-01-01", "2026-08-01")
|
||||
assert window.end == "2026-08-01"
|
||||
assert clamp_window(date(2020, 1, 1), END).end == "2026-08-01"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -189,7 +191,7 @@ class TestWarmingEstimate:
|
|||
class TestReadPathSource:
|
||||
async def test_the_dashboard_query_never_touches_the_spend_logs(self):
|
||||
prisma = _FakePrisma([group_row()])
|
||||
await compute_benchmarks(prisma, GROUP_KINDS, "2026-07-02", "2026-08-01")
|
||||
await compute_benchmarks(prisma, GROUP_KINDS, START, END)
|
||||
sql = prisma.db.queries[0][0]
|
||||
assert "LiteLLM_SpendLogs" not in sql
|
||||
assert "LiteLLM_AutoRouterSession" in sql
|
||||
|
|
@ -197,14 +199,14 @@ class TestReadPathSource:
|
|||
async def test_one_query_covers_every_configured_auto_router(self):
|
||||
prisma = _FakePrisma([group_row(), group_row(model_group="claude-router-2")])
|
||||
result = await compute_benchmarks(
|
||||
prisma, {"claude-auto": "semantic", "claude-router-2": "complexity"}, "2026-07-02", "2026-08-01"
|
||||
prisma, {"claude-auto": "semantic", "claude-router-2": "complexity"}, START, END
|
||||
)
|
||||
assert len(prisma.db.queries) == 1
|
||||
assert {g.model_group for g in result.groups} == {"claude-auto", "claude-router-2"}
|
||||
|
||||
async def test_each_group_is_labelled_with_its_router_kind(self):
|
||||
prisma = _FakePrisma([group_row(model_group="claude-router-2")])
|
||||
result = await compute_benchmarks(prisma, {"claude-router-2": "complexity"}, "2026-07-02", "2026-08-01")
|
||||
result = await compute_benchmarks(prisma, {"claude-router-2": "complexity"}, START, END)
|
||||
assert result.groups[0].router_kind == "complexity"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.proxy.spend_tracking.auto_router_sessions import (
|
|||
PROMPT_CACHE_TTL_SECONDS,
|
||||
SessionState,
|
||||
TurnFacts,
|
||||
fold_session,
|
||||
fold_turn,
|
||||
merge_deltas,
|
||||
state_from_row,
|
||||
|
|
@ -201,6 +202,36 @@ class TestTtlSelection:
|
|||
assert inside_the_hour.stale_return_misses == 0
|
||||
|
||||
|
||||
class TestFoldSession:
|
||||
"""An interval's turns are folded in start order, not in arrival order."""
|
||||
|
||||
def test_turns_that_arrived_out_of_order_still_all_classify(self):
|
||||
"""Two turns in flight together finish in whichever order the providers answer."""
|
||||
arrival_order = (turn(MODEL_B, at=60, created=5000), turn(MODEL_A, at=0, created=5000))
|
||||
folded = fold_session(EMPTY_SESSION_STATE, arrival_order, rates=rates)
|
||||
assert folded.turns == 2
|
||||
assert buckets(folded) == 2
|
||||
assert folded.state.last_model == MODEL_B
|
||||
|
||||
def test_folding_an_interval_matches_folding_its_turns_one_at_a_time(self):
|
||||
turns = (
|
||||
turn(MODEL_A, at=0, created=5000),
|
||||
turn(MODEL_A, at=60, read=5000),
|
||||
turn(MODEL_B, at=400, created=5000),
|
||||
turn(MODEL_A, at=800, created=5000),
|
||||
)
|
||||
deltas, state = fold_all(turns)
|
||||
folded = fold_session(EMPTY_SESSION_STATE, turns, rates=rates)
|
||||
assert folded.turns == sum(d.turns for d in deltas)
|
||||
assert folded.return_turns == sum(d.return_turns for d in deltas)
|
||||
assert folded.replay_spend == pytest.approx(sum(d.replay_spend for d in deltas))
|
||||
assert folded.state == state
|
||||
|
||||
def test_an_empty_interval_folds_to_the_state_it_was_given(self):
|
||||
_, state = fold_all((turn(MODEL_A, at=0, created=5000),))
|
||||
assert fold_session(state, (), rates=rates).state is state
|
||||
|
||||
|
||||
class TestOutOfOrderTurns:
|
||||
def test_a_late_turn_keeps_its_spend_but_not_its_classification(self):
|
||||
first = fold_turn(EMPTY_SESSION_STATE, turn(MODEL_A, at=100, created=5000, spend=0.02), rates=rates)
|
||||
|
|
@ -336,18 +367,33 @@ def test_state_from_row_without_a_timestamp_starts_at_the_epoch():
|
|||
)
|
||||
|
||||
|
||||
class _RecordingTable:
|
||||
"""A session-rollup table that can be told to fail."""
|
||||
class _StoredRow:
|
||||
"""A session rollup as prisma hands it back."""
|
||||
|
||||
def __init__(self, fail: bool = False):
|
||||
self.fail = fail
|
||||
def __init__(self, last_model: str, last_turn_at: float, model_state: dict):
|
||||
self.last_model = last_model
|
||||
self.last_turn_at = datetime.fromtimestamp(last_turn_at, tz=timezone.utc)
|
||||
self.model_state = model_state
|
||||
|
||||
|
||||
class _RecordingTable:
|
||||
"""A session-rollup table that can be told to fail on either side."""
|
||||
|
||||
def __init__(self, fail_read: bool = False, fail_write: bool = False, row=None):
|
||||
self.fail_read = fail_read
|
||||
self.fail_write = fail_write
|
||||
self.row = row
|
||||
self.reads = 0
|
||||
self.upserts: list = []
|
||||
|
||||
async def find_unique(self, where):
|
||||
return None
|
||||
self.reads += 1
|
||||
if self.fail_read:
|
||||
raise RuntimeError("transient database fault")
|
||||
return self.row
|
||||
|
||||
async def upsert(self, where, data):
|
||||
if self.fail:
|
||||
if self.fail_write:
|
||||
raise RuntimeError("transient database fault")
|
||||
self.upserts.append((where, data))
|
||||
|
||||
|
|
@ -361,74 +407,156 @@ def _turn_at(at: float, model: str = MODEL_A) -> TurnFacts:
|
|||
return turn(model, at=at, created=5000)
|
||||
|
||||
|
||||
def _stored_session() -> _StoredRow:
|
||||
"""A session last served on MODEL_B that has already used MODEL_A."""
|
||||
return _StoredRow(
|
||||
last_model=MODEL_B,
|
||||
last_turn_at=60.0,
|
||||
model_state={
|
||||
MODEL_A: {"last_used_at": 0.0, "provisioned_replay_spend": 0.0},
|
||||
MODEL_B: {"last_used_at": 60.0, "provisioned_replay_spend": 0.0},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestTheLoggingPathNeverTouchesTheDatabase:
|
||||
"""Staging a turn must not put a database round trip in front of spend tracking."""
|
||||
|
||||
async def test_turns_can_be_recorded_with_no_database_in_sight(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
queue = AutoRouterSessionQueue()
|
||||
for at in (0, 60, 120):
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at))
|
||||
|
||||
table = _RecordingTable()
|
||||
assert await queue.flush(_RecordingPrisma(table)) == 1
|
||||
assert table.upserts[0][1]["create"]["turns"] == 3
|
||||
|
||||
async def test_a_session_costs_one_read_per_flush_at_most_not_one_per_turn(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable()
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
for at in (0, 60, 120):
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at))
|
||||
await queue.flush(prisma)
|
||||
|
||||
assert table.reads == 1
|
||||
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(180))
|
||||
await queue.flush(prisma)
|
||||
assert table.reads == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestFlushDurability:
|
||||
"""A transient write fault must not silently delete an interval of traffic."""
|
||||
"""A transient database fault must not delete an interval of traffic or its history."""
|
||||
|
||||
async def test_a_failed_write_is_restaged_rather_than_dropped(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable(fail=True)
|
||||
table = _RecordingTable(fail_write=True)
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma)
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0))
|
||||
|
||||
assert await queue.flush(prisma) == 0
|
||||
|
||||
table.fail = False
|
||||
table.fail_write = False
|
||||
assert await queue.flush(prisma) == 1
|
||||
assert table.upserts[0][1]["create"]["turns"] == 1
|
||||
|
||||
async def test_a_restaged_batch_merges_under_turns_staged_since(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable(fail=True)
|
||||
table = _RecordingTable(fail_write=True)
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma)
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0))
|
||||
await queue.flush(prisma)
|
||||
|
||||
table.fail = False
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60), prisma)
|
||||
table.fail_write = False
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60))
|
||||
assert await queue.flush(prisma) == 1
|
||||
# Both turns land, once each
|
||||
assert table.upserts[0][1]["create"]["turns"] == 2
|
||||
|
||||
async def test_a_failed_state_read_writes_nothing_at_all(self):
|
||||
"""The corrupting move was folding onto an empty state and persisting it."""
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable(fail_read=True, row=_stored_session())
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120))
|
||||
|
||||
assert await queue.flush(prisma) == 0
|
||||
assert table.upserts == []
|
||||
|
||||
async def test_history_survives_a_failed_read_and_still_classifies_the_turn(self):
|
||||
"""Returning to a tier this session already used is a return, not a first visit."""
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable(fail_read=True, row=_stored_session())
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120))
|
||||
await queue.flush(prisma)
|
||||
|
||||
table.fail_read = False
|
||||
assert await queue.flush(prisma) == 1
|
||||
update = table.upserts[0][1]["update"]
|
||||
assert update["return_turns"] == {"increment": 1}
|
||||
assert update["first_visit_turns"] == {"increment": 0}
|
||||
assert update["turns"] == {"increment": 1}
|
||||
|
||||
async def test_a_successful_flush_stages_nothing_back(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
prisma = _RecordingPrisma(_RecordingTable())
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma)
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0))
|
||||
assert await queue.flush(prisma) == 1
|
||||
assert await queue.flush(prisma) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestPendingIsBounded:
|
||||
"""`session_id` is caller-controlled, so the staged aggregate needs a ceiling."""
|
||||
class TestStagingIsBounded:
|
||||
"""`session_id` is caller-controlled, so what is held between flushes needs a ceiling."""
|
||||
|
||||
async def test_new_sessions_are_refused_once_the_aggregate_is_full(self):
|
||||
async def test_turns_are_refused_once_the_staging_is_full(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
prisma = _RecordingPrisma(_RecordingTable())
|
||||
queue = AutoRouterSessionQueue(max_tracked_sessions=2)
|
||||
table = _RecordingTable()
|
||||
queue = AutoRouterSessionQueue(max_staged_turns=2)
|
||||
for i in range(5):
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0), prisma)
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0))
|
||||
|
||||
assert await queue.flush(prisma) == 2
|
||||
assert await queue.flush(_RecordingPrisma(table)) == 2
|
||||
|
||||
async def test_a_session_already_staged_keeps_accumulating_at_the_cap(self):
|
||||
"""Refusing new keys must not stall the conversations already in flight."""
|
||||
async def test_the_ceiling_counts_turns_not_sessions(self):
|
||||
"""One caller replaying a long session must not evade the memory bound."""
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable()
|
||||
queue = AutoRouterSessionQueue(max_staged_turns=2)
|
||||
for at in (0, 60, 120, 180):
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(at))
|
||||
|
||||
await queue.flush(_RecordingPrisma(table))
|
||||
assert table.upserts[0][1]["create"]["turns"] == 2
|
||||
|
||||
async def test_staging_reopens_once_it_has_drained(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable()
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue(max_tracked_sessions=1)
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0), prisma)
|
||||
await queue.record_turn(("s2", "g"), "complexity", None, _turn_at(0), prisma)
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60), prisma)
|
||||
|
||||
queue = AutoRouterSessionQueue(max_staged_turns=1)
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(0))
|
||||
await queue.flush(prisma)
|
||||
assert len(table.upserts) == 1
|
||||
assert table.upserts[0][1]["create"]["turns"] == 2
|
||||
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(60))
|
||||
assert await queue.flush(prisma) == 1
|
||||
assert table.upserts[1][1]["update"]["turns"] == {"increment": 1}
|
||||
|
|
|
|||
|
|
@ -10825,3 +10825,50 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog):
|
|||
ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={})
|
||||
|
||||
assert MOCK_TESTING_CONFIG_KEY not in caplog.text
|
||||
|
||||
|
||||
def _auto_router_benchmarks_client(monkeypatch):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
monkeypatch.setattr(ps, "llm_router", None)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"params, expected_status",
|
||||
[
|
||||
({"start_date": "not-a-date", "end_date": "2026-08-01"}, 422),
|
||||
({"start_date": "2026-08-01", "end_date": "2026-13-45"}, 422),
|
||||
({"start_date": "2026-08-01", "end_date": "2026-07-02"}, 400),
|
||||
],
|
||||
)
|
||||
def test_auto_router_benchmarks_rejects_windows_it_cannot_serve(
|
||||
monkeypatch, params, expected_status
|
||||
):
|
||||
"""A malformed date used to surface as a 500 and an inverted range as an empty
|
||||
dashboard, which reads to an operator as "no auto-router traffic" rather than as a
|
||||
bad request. Both are the caller's error and must say so. 404 (no auto-router
|
||||
configured) is checked last in the handler, so reaching it means the window passed."""
|
||||
client = _auto_router_benchmarks_client(monkeypatch)
|
||||
try:
|
||||
resp = client.get("/auto_router/benchmarks", params=params)
|
||||
assert resp.status_code == expected_status, resp.text
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_auto_router_benchmarks_accepts_a_well_formed_window(monkeypatch):
|
||||
"""The guard must not reject the windows the dashboard actually sends."""
|
||||
client = _auto_router_benchmarks_client(monkeypatch)
|
||||
try:
|
||||
resp = client.get(
|
||||
"/auto_router/benchmarks",
|
||||
params={"start_date": "2026-07-02", "end_date": "2026-08-01"},
|
||||
)
|
||||
assert resp.status_code == 404, resp.text
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
|
|
|||
9
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
9
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -776,10 +776,11 @@ export interface paths {
|
|||
* counterfactual baseline, and how the provider prompt cache behaved.
|
||||
*
|
||||
* Reads the per-session rollup, never the per-request spend logs.
|
||||
* ``start_date`` / ``end_date`` are ``YYYY-MM-DD``; the window is clamped to
|
||||
* the most recent ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes
|
||||
* the window actually served. Sessions are counted whole when they were active
|
||||
* in the window. Returns 404 when no auto-router is configured.
|
||||
* ``start_date`` / ``end_date`` are ``YYYY-MM-DD``, rejected by the framework
|
||||
* when malformed; the window is clamped to the most recent
|
||||
* ``BENCHMARKS_MAX_WINDOW_DAYS`` days and the response echoes the window
|
||||
* actually served. Sessions are counted whole when they were active in the
|
||||
* window. Returns 404 when no auto-router is configured.
|
||||
*/
|
||||
get: operations["get_auto_router_benchmarks_auto_router_benchmarks_get"];
|
||||
put?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue