mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
perf(spend): drain auto-router session rollups a chunk at a time
The flusher read and wrote one session per round trip, so an interval that staged N sessions cost 2N sequential statements. It now drains in chunks of at most a thousand sessions, each one a single find_many over the composite keys this pod has not already cached followed by a single transaction carrying every upsert, which holds the flush at a handful of round trips no matter how many sessions an interval touched The chunk is also the unit of failure. Reading, folding and writing are still one unit of work, so a chunk that faults at either end writes nothing and re-stages every session in it, and a replay cannot double-count. The state cache still advances only once the write has landed, an unreadable chunk is still never folded as empty history, and record_turn still does no I/O at all
This commit is contained in:
Tin Chi Lo
parent
8d300d98d8
commit
da356f86d3
3 changed files with 272 additions and 106 deletions
|
|
@ -9,20 +9,25 @@ 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.
|
||||
|
||||
That unit is a chunk of sessions rather than one session, so an interval costs a
|
||||
handful of round trips instead of two per session: one ``find_many`` for the
|
||||
states this pod has not cached, then one transaction carrying every upsert.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.spend_tracking.auto_router_sessions import (
|
||||
EMPTY_SESSION_STATE,
|
||||
SessionState,
|
||||
StateLookup,
|
||||
StateUnavailable,
|
||||
TurnDelta,
|
||||
TurnFacts,
|
||||
|
|
@ -40,6 +45,7 @@ SessionKey = tuple[str, str]
|
|||
|
||||
DEFAULT_MAX_STAGED_TURNS = 50_000
|
||||
MAX_CACHED_SESSIONS = 10_000
|
||||
SESSIONS_PER_STATEMENT = 1_000
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -49,6 +55,11 @@ class _Pending:
|
|||
turns: tuple[TurnFacts, ...]
|
||||
|
||||
|
||||
_Chunk = Mapping[SessionKey, _Pending]
|
||||
_Folded = tuple[SessionKey, _Pending, TurnDelta]
|
||||
_ChunkStates = Mapping[SessionKey, SessionState] | StateUnavailable
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _warn_staging_full(cap: int) -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -62,6 +73,56 @@ def _epoch_to_datetime(value: float) -> datetime:
|
|||
return datetime.fromtimestamp(value, tz=timezone.utc)
|
||||
|
||||
|
||||
def _chunked(batch: _Chunk, size: int) -> tuple[_Chunk, ...]:
|
||||
"""``batch`` split into runs of at most ``size``, each in key order.
|
||||
|
||||
Sorting once here is what gives every pod the same lock ordering, and the
|
||||
size cap is what stops one statement from carrying the whole staging area.
|
||||
"""
|
||||
keys = sorted(batch)
|
||||
return tuple(
|
||||
MappingProxyType({key: batch[key] for key in keys[at : at + size]}) for at in range(0, len(keys), size)
|
||||
)
|
||||
|
||||
|
||||
def _key_fields(key: SessionKey) -> Mapping[str, str]:
|
||||
"""One composite primary key as the two columns that make it up."""
|
||||
session_id, model_group = key
|
||||
return {"session_id": session_id, "model_group": model_group} # mutable-ok: prisma's query API takes dict payloads
|
||||
|
||||
|
||||
def _unique_where(key: SessionKey) -> Mapping[str, Mapping[str, str]]:
|
||||
"""The composite primary key, under the name prisma gives the ``@@id`` selector."""
|
||||
return {"session_id_model_group": _key_fields(key)} # mutable-ok: prisma's query API takes dict payloads
|
||||
|
||||
|
||||
def _increment(value: float) -> Mapping[str, float]:
|
||||
"""One counter's atomic add, the way prisma spells it in an update payload."""
|
||||
return {"increment": value} # mutable-ok: prisma's write API takes dict payloads
|
||||
|
||||
|
||||
def _upsert_data(key: SessionKey, pending: _Pending, delta: TurnDelta) -> Mapping[str, Mapping[str, object]]:
|
||||
"""One session's write: create its row, or add this interval onto the row already there."""
|
||||
counters = counters_of(delta)
|
||||
increments = MappingProxyType({name: _increment(value) for name, value in counters.items()})
|
||||
shared = { # mutable-ok: prisma's write API takes dict payloads
|
||||
"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,
|
||||
}
|
||||
return { # mutable-ok: prisma's write API takes dict payloads
|
||||
"create": { # mutable-ok: prisma's write API takes dict payloads
|
||||
**_key_fields(key),
|
||||
"router_kind": pending.router_kind,
|
||||
"first_turn_at": _epoch_to_datetime(min(turn.started_at for turn in pending.turns)),
|
||||
**shared,
|
||||
**counters,
|
||||
},
|
||||
"update": {**increments, **shared}, # mutable-ok: prisma's write API takes dict payloads
|
||||
}
|
||||
|
||||
|
||||
class AutoRouterSessionQueue:
|
||||
"""Stages auto-routed turns in memory and folds them into the rollup on flush."""
|
||||
|
||||
|
|
@ -72,13 +133,7 @@ class AutoRouterSessionQueue:
|
|||
self._lock = asyncio.Lock()
|
||||
self._max_staged_turns = max_staged_turns
|
||||
|
||||
async def record_turn(
|
||||
self,
|
||||
key: SessionKey,
|
||||
router_kind: str,
|
||||
baseline_model: str | None,
|
||||
turn: TurnFacts,
|
||||
) -> None:
|
||||
async def record_turn(self, key: SessionKey, router_kind: str, baseline_model: str | None, turn: TurnFacts) -> None:
|
||||
"""Stage one turn against its session. Does no I/O; the fold happens on flush.
|
||||
|
||||
``session_id`` is caller-controlled, so the staging is capped on turns
|
||||
|
|
@ -111,20 +166,19 @@ class AutoRouterSessionQueue:
|
|||
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.
|
||||
A chunk that could not be read or written is staged again rather than
|
||||
dropped. Reading, folding and writing are one unit per chunk, so a
|
||||
failure means nothing in it 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:
|
||||
batch = self._pending
|
||||
self._pending = {} # mutable-ok: fresh staging for the next interval
|
||||
self._staged_turns = 0
|
||||
|
||||
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)
|
||||
}
|
||||
chunks = _chunked(batch, SESSIONS_PER_STATEMENT)
|
||||
failed = tuple([key for chunk in chunks if not await self._commit(chunk, prisma_client) for key in chunk])
|
||||
if failed:
|
||||
verbose_proxy_logger.warning(
|
||||
"auto_router_sessions: %d of %d sessions could not be folded; re-staging them for the next flush",
|
||||
|
|
@ -132,96 +186,80 @@ class AutoRouterSessionQueue:
|
|||
len(batch),
|
||||
)
|
||||
async with self._lock:
|
||||
for key, pending in failed.items():
|
||||
self._stage(key, pending.router_kind, pending.baseline_model, pending.turns)
|
||||
for key in failed:
|
||||
self._stage(key, batch[key].router_kind, batch[key].baseline_model, batch[key].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.
|
||||
async def _commit(self, chunk: _Chunk, prisma_client: "PrismaClient") -> bool:
|
||||
"""Read a chunk's states, fold its staged turns onto them, write them all 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.
|
||||
The chunk lands whole or not at all: one read covers it and one
|
||||
transaction writes it, so a failure at either end leaves nothing written
|
||||
for any session in it and the caller re-stages them all. 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 and folds identically.
|
||||
"""
|
||||
state = await self._session_state(key, prisma_client)
|
||||
if isinstance(state, StateUnavailable):
|
||||
states = await self._session_states(tuple(chunk), prisma_client)
|
||||
if isinstance(states, StateUnavailable):
|
||||
return False
|
||||
delta = fold_session(state, pending.turns)
|
||||
if not await self._write(key, pending, delta, prisma_client):
|
||||
folded = tuple((key, pending, fold_session(states[key], pending.turns)) for key, pending in chunk.items())
|
||||
if not await self._write(folded, prisma_client):
|
||||
return False
|
||||
self._state[key] = delta.state
|
||||
self._state.move_to_end(key)
|
||||
for key, _, delta in folded:
|
||||
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.
|
||||
async def _session_states(self, keys: tuple[SessionKey, ...], prisma_client: "PrismaClient") -> _ChunkStates:
|
||||
"""Each key's state: from memory where this pod has folded it before, else from one read.
|
||||
|
||||
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.
|
||||
The filter is an OR over whole composite keys rather than ``session_id IN
|
||||
(...) AND model_group IN (...)``, because the latter is a cross product:
|
||||
it would drag back rows for pairs nobody asked about, growing with the
|
||||
number of auto-routers in the flush and leaving Python to discard them.
|
||||
|
||||
Only the flusher touches this cache, so it needs no lock, and a chunk it
|
||||
has entirely cached costs no read at all. A fault leaves the whole chunk
|
||||
unavailable, which re-stages it; folding onto an empty state instead
|
||||
would replace a history that really happened with one derived from a
|
||||
single interval.
|
||||
"""
|
||||
cached = self._state.get(key)
|
||||
if cached is not None:
|
||||
self._state.move_to_end(key)
|
||||
return cached
|
||||
session_id, model_group = key
|
||||
missing = tuple(key for key in keys if key not in self._state)
|
||||
if not missing:
|
||||
return MappingProxyType({key: self._state[key] for key in keys})
|
||||
try:
|
||||
row = await AutoRouterSessionRepository(prisma_client).table.find_unique(
|
||||
where={ # mutable-ok: prisma's write API takes dict payloads
|
||||
"session_id_model_group": { # mutable-ok: a JSON object is a dict by definition
|
||||
"session_id": session_id,
|
||||
"model_group": model_group,
|
||||
}
|
||||
}
|
||||
rows = await AutoRouterSessionRepository(prisma_client).table.find_many(
|
||||
where={"OR": [_key_fields(key) for key in missing]} # mutable-ok: prisma's query API takes dicts
|
||||
)
|
||||
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)
|
||||
except Exception as e: # noqa: BLE001 # a read fault re-stages the chunk rather than failing the flush
|
||||
verbose_proxy_logger.warning("auto_router_sessions: could not load %d session states (%s)", len(missing), 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)
|
||||
|
||||
async def _write(self, key: SessionKey, pending: _Pending, delta: TurnDelta, prisma_client: "PrismaClient") -> bool:
|
||||
session_id, model_group = key
|
||||
counters = counters_of(delta)
|
||||
shared = { # mutable-ok: prisma's write API takes dict payloads
|
||||
"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,
|
||||
stored = { # mutable-ok: built once from the rows this read returned
|
||||
(row.session_id, row.model_group): state_from_row(row.last_model, row.last_turn_at, row.model_state)
|
||||
for row in rows
|
||||
}
|
||||
return MappingProxyType(
|
||||
{key: self._state[key] if key in self._state else stored.get(key, EMPTY_SESSION_STATE) for key in keys}
|
||||
)
|
||||
|
||||
async def _write(self, folded: tuple[_Folded, ...], prisma_client: "PrismaClient") -> bool:
|
||||
"""Write a chunk's rows as one transaction of atomic increment upserts.
|
||||
|
||||
``batch_()`` issues its statements sequentially inside the transaction, so
|
||||
iteration order is lock acquisition order; the chunk arrives sorted, which
|
||||
is what keeps two pods flushing the same sessions from deadlocking.
|
||||
"""
|
||||
try:
|
||||
await AutoRouterSessionRepository(
|
||||
prisma_client
|
||||
).table.upsert(
|
||||
where={ # mutable-ok: prisma's write API takes dict payloads
|
||||
"session_id_model_group": { # mutable-ok: a JSON object is a dict by definition
|
||||
"session_id": session_id,
|
||||
"model_group": model_group,
|
||||
}
|
||||
},
|
||||
data={ # mutable-ok: prisma's write API takes dict payloads
|
||||
"create": { # mutable-ok: prisma's write API takes dict payloads
|
||||
"session_id": session_id,
|
||||
"model_group": model_group,
|
||||
"router_kind": pending.router_kind,
|
||||
"first_turn_at": _epoch_to_datetime(min(turn.started_at for turn in pending.turns)),
|
||||
**shared,
|
||||
**counters,
|
||||
},
|
||||
"update": { # mutable-ok: prisma's write API takes dict payloads
|
||||
**{ # mutable-ok: a JSON object is a dict by definition
|
||||
field: {"increment": value} # mutable-ok: a JSON object is a dict by definition
|
||||
for field, value in counters.items() # mutable-ok: spread into the prisma payload immediately below
|
||||
}, # mutable-ok: spread into the prisma payload immediately below
|
||||
**shared,
|
||||
},
|
||||
},
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # one session's write must not drop the rest of the batch
|
||||
verbose_proxy_logger.exception("auto_router_sessions: failed to flush session %s (%s)", key, e)
|
||||
async with prisma_client.db.tx(timeout=timedelta(seconds=60)) as transaction:
|
||||
async with transaction.batch_() as batcher:
|
||||
table = batcher.litellm_autoroutersession
|
||||
for key, pending, delta in folded:
|
||||
table.upsert(where=_unique_where(key), data=_upsert_data(key, pending, delta))
|
||||
except Exception as e: # noqa: BLE001 # one chunk's write must not drop the rest of the flush
|
||||
verbose_proxy_logger.exception("auto_router_sessions: failed to flush %d sessions (%s)", len(folded), e)
|
||||
return False
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -91,9 +91,6 @@ class StateUnavailable:
|
|||
"""
|
||||
|
||||
|
||||
StateLookup = SessionState | StateUnavailable
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TurnFacts:
|
||||
"""One auto-routed request, as the spend writer sees it.
|
||||
|
|
|
|||
|
|
@ -370,46 +370,103 @@ def test_state_from_row_without_a_timestamp_starts_at_the_epoch():
|
|||
class _StoredRow:
|
||||
"""A session rollup as prisma hands it back."""
|
||||
|
||||
def __init__(self, last_model: str, last_turn_at: float, model_state: dict):
|
||||
def __init__(self, session_id: str, model_group: str, last_model: str, last_turn_at: float, model_state: dict):
|
||||
self.session_id = session_id
|
||||
self.model_group = model_group
|
||||
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."""
|
||||
"""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):
|
||||
Counts round trips the way prisma issues them: one `find_many` however many
|
||||
keys it selects, and one transaction however many upserts it carries.
|
||||
"""
|
||||
|
||||
def __init__(self, fail_read: bool = False, fail_write: bool = False, rows=()):
|
||||
self.fail_read = fail_read
|
||||
self.fail_write = fail_write
|
||||
self.row = row
|
||||
self.rows = rows
|
||||
self.reads = 0
|
||||
self.transactions = 0
|
||||
self.upserts: list = []
|
||||
|
||||
async def find_unique(self, where):
|
||||
async def find_many(self, where):
|
||||
self.reads += 1
|
||||
if self.fail_read:
|
||||
raise RuntimeError("transient database fault")
|
||||
return self.row
|
||||
wanted = {(pair["session_id"], pair["model_group"]) for pair in where["OR"]}
|
||||
return [row for row in self.rows if (row.session_id, row.model_group) in wanted]
|
||||
|
||||
async def upsert(self, where, data):
|
||||
def commit(self, statements):
|
||||
self.transactions += 1
|
||||
if self.fail_write:
|
||||
raise RuntimeError("transient database fault")
|
||||
self.upserts.append((where, data))
|
||||
self.upserts.extend(statements)
|
||||
|
||||
|
||||
class _RecordingBatchActions:
|
||||
"""Prisma's batcher: statements queue up and land only when the batch commits."""
|
||||
|
||||
def __init__(self):
|
||||
self.queued: list = []
|
||||
|
||||
def upsert(self, where, data):
|
||||
self.queued.append((where, data))
|
||||
|
||||
|
||||
class _RecordingBatch:
|
||||
def __init__(self, table):
|
||||
self._table = table
|
||||
self.litellm_autoroutersession = _RecordingBatchActions()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if exc is None:
|
||||
self._table.commit(tuple(self.litellm_autoroutersession.queued))
|
||||
return False
|
||||
|
||||
|
||||
class _RecordingTransaction:
|
||||
def __init__(self, table):
|
||||
self._table = table
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def batch_(self):
|
||||
return _RecordingBatch(self._table)
|
||||
|
||||
|
||||
class _RecordingDb:
|
||||
def __init__(self, table):
|
||||
self.litellm_autoroutersession = table
|
||||
|
||||
def tx(self, timeout=None):
|
||||
return _RecordingTransaction(self.litellm_autoroutersession)
|
||||
|
||||
|
||||
class _RecordingPrisma:
|
||||
def __init__(self, table):
|
||||
self.db = type("_Db", (), {"litellm_autoroutersession": table})()
|
||||
self.db = _RecordingDb(table)
|
||||
|
||||
|
||||
def _turn_at(at: float, model: str = MODEL_A) -> TurnFacts:
|
||||
return turn(model, at=at, created=5000)
|
||||
|
||||
|
||||
def _stored_session() -> _StoredRow:
|
||||
def _stored_session(session_id: str = "s1", model_group: str = "g") -> _StoredRow:
|
||||
"""A session last served on MODEL_B that has already used MODEL_A."""
|
||||
return _StoredRow(
|
||||
session_id=session_id,
|
||||
model_group=model_group,
|
||||
last_model=MODEL_B,
|
||||
last_turn_at=60.0,
|
||||
model_state={
|
||||
|
|
@ -487,7 +544,7 @@ class TestFlushDurability:
|
|||
"""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())
|
||||
table = _RecordingTable(fail_read=True, rows=(_stored_session(),))
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120))
|
||||
|
|
@ -499,7 +556,7 @@ class TestFlushDurability:
|
|||
"""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())
|
||||
table = _RecordingTable(fail_read=True, rows=(_stored_session(),))
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
await queue.record_turn(("s1", "g"), "complexity", None, _turn_at(120))
|
||||
|
|
@ -560,3 +617,77 @@ class TestStagingIsBounded:
|
|||
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}
|
||||
|
||||
|
||||
def test_chunking_caps_what_one_statement_carries_and_keeps_it_in_key_order():
|
||||
"""Key order is the lock order two pods draining the same sessions have to agree on."""
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import SESSIONS_PER_STATEMENT, _chunked
|
||||
|
||||
batch = {(f"s{i:05d}", "g"): i for i in range(2500)}
|
||||
chunks = _chunked(batch, SESSIONS_PER_STATEMENT)
|
||||
|
||||
assert [len(chunk) for chunk in chunks] == [1000, 1000, 500]
|
||||
assert [key for chunk in chunks for key in chunk] == sorted(batch)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestFlushCostDoesNotGrowWithSessionCount:
|
||||
"""A read and an upsert per session is two round trips per session; a flush must not do that."""
|
||||
|
||||
async def test_many_sessions_cost_one_read_and_one_transaction(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable()
|
||||
queue = AutoRouterSessionQueue()
|
||||
for i in range(250):
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0))
|
||||
|
||||
assert await queue.flush(_RecordingPrisma(table)) == 250
|
||||
assert (table.reads, table.transactions) == (1, 1)
|
||||
assert len(table.upserts) == 250
|
||||
|
||||
async def test_sessions_this_pod_has_already_folded_cost_no_read_at_all(self):
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable()
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
for i in range(50):
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0))
|
||||
await queue.flush(prisma)
|
||||
|
||||
for i in range(50):
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(60))
|
||||
assert await queue.flush(prisma) == 50
|
||||
assert (table.reads, table.transactions) == (1, 2)
|
||||
|
||||
async def test_a_chunk_that_could_not_be_written_restages_every_session_in_it(self):
|
||||
"""The transaction is the failure unit, so nothing in it landed and nothing may be dropped."""
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable(fail_write=True)
|
||||
prisma = _RecordingPrisma(table)
|
||||
queue = AutoRouterSessionQueue()
|
||||
for i in range(5):
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(0))
|
||||
|
||||
assert await queue.flush(prisma) == 0
|
||||
assert table.upserts == []
|
||||
|
||||
table.fail_write = False
|
||||
assert await queue.flush(prisma) == 5
|
||||
written = {where["session_id_model_group"]["session_id"] for where, _ in table.upserts}
|
||||
assert written == {f"s{i}" for i in range(5)}
|
||||
|
||||
async def test_one_read_serves_a_chunk_of_sessions_that_all_have_history(self):
|
||||
"""Every session's own row has to come back from the batched read, not just the first."""
|
||||
from litellm.proxy.spend_tracking.auto_router_session_queue import AutoRouterSessionQueue
|
||||
|
||||
table = _RecordingTable(rows=tuple(_stored_session(session_id=f"s{i}") for i in range(5)))
|
||||
queue = AutoRouterSessionQueue()
|
||||
for i in range(5):
|
||||
await queue.record_turn((f"s{i}", "g"), "complexity", None, _turn_at(120))
|
||||
|
||||
assert await queue.flush(_RecordingPrisma(table)) == 5
|
||||
assert table.reads == 1
|
||||
assert all(data["update"]["return_turns"] == {"increment": 1} for _, data in table.upserts)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue