mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
perf(proxy): collapse per-worker SGR upserts into one statement per flush (#40362)
Each proxy worker flushed one Prisma upsert per active (date, category, route) bucket every interval, so the Postgres primary saw workers x routes statements per interval across the deployment. A flush now builds a single multi-row INSERT ... ON CONFLICT DO UPDATE, and with use_redis_transaction_buffer on the workers push snapshots to a Redis list that one lease-holding pod folds and commits, so the whole deployment costs one statement per interval. The leader keeps popping until the list is empty so a deployment wider than the dequeue cap cannot build a backlog, and rows that fail both the commit and the Redis re-queue fall back to the leader's own accumulator instead of being lost. Resolves LIT-7371 Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
699ae63b2a
commit
c7163a80dd
4 changed files with 537 additions and 89 deletions
|
|
@ -335,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv(
|
|||
|
||||
########### v2 Architecture constants for managing writing updates to the database ###########
|
||||
REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer"
|
||||
REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer"
|
||||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer"
|
||||
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer"
|
||||
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer"
|
||||
|
|
|
|||
|
|
@ -10,13 +10,28 @@ strings rather than passing the raw path through. Nothing a caller sends can
|
|||
add a key, so the fold and the table it commits to are bounded by (days x
|
||||
routes) however much traffic arrives, and the response path carries no
|
||||
unbounded queue that would block once full.
|
||||
|
||||
A flush commits its whole snapshot as one multi-row ``INSERT ... ON CONFLICT DO
|
||||
UPDATE`` rather than one upsert per key, so a worker costs the primary one
|
||||
statement per interval however many routes it served. With
|
||||
``use_redis_transaction_buffer`` on, workers instead push their snapshot to a
|
||||
Redis list and one lock-holding pod folds every entry and writes the table, so
|
||||
the deployment as a whole costs the primary one statement per interval.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Iterable
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import RedisCache
|
||||
from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.middleware.billable_request_metrics_middleware import BillableCategory
|
||||
from litellm.types.proxy.gateway_requests import (
|
||||
GatewayRequestCounts,
|
||||
|
|
@ -28,6 +43,15 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_EMPTY: Final = GatewayRequestCounts(successful_requests=0, failed_requests=0)
|
||||
_TABLE: Final = '"LiteLLM_DailyGatewayRequests"'
|
||||
_COLUMNS_PER_ROW: Final = 5
|
||||
_UTC_NOW: Final = "(NOW() AT TIME ZONE 'UTC')"
|
||||
GATEWAY_REQUESTS_JOB_NAME: Final = "update_gateway_requests_job"
|
||||
|
||||
_BufferedRows: TypeAlias = tuple[tuple[str, str, str, int, int], ...]
|
||||
_BUFFERED_ROWS: Final = TypeAdapter(_BufferedRows)
|
||||
_BUFFERED_ENTRIES: Final = TypeAdapter(tuple[str | bytes, ...])
|
||||
_NO_COUNTS: Final[GatewayRequestSnapshot] = MappingProxyType({})
|
||||
|
||||
|
||||
def _utc_date() -> str:
|
||||
|
|
@ -59,20 +83,54 @@ class GatewayRequestAccumulator:
|
|||
route) however long the database is unreachable.
|
||||
|
||||
This buys at-least-once, not exactly-once, and the cost is worth stating.
|
||||
The batch commits inside its context manager's ``__aexit__``, so a failure
|
||||
raised after the transaction committed (a connection dropped while reading
|
||||
the acknowledgement) restores counts that are already persisted, and the
|
||||
next flush increments them a second time. Exactly-once would need a dedup
|
||||
key the upserts could ignore on replay. For a traffic-volume metric a rare
|
||||
The statement commits on the server before its acknowledgement is read, so
|
||||
a failure raised after the commit (a connection dropped while reading the
|
||||
acknowledgement) restores counts that are already persisted, and the next
|
||||
flush increments them a second time. Exactly-once would need a dedup key
|
||||
the upsert could ignore on replay. For a traffic-volume metric a rare
|
||||
overcount on a dropped acknowledgement beats losing a whole interval to
|
||||
every database blip, so the trade is deliberate.
|
||||
"""
|
||||
for key, counts in snapshot.items():
|
||||
existing = self._counts.get(key, _EMPTY)
|
||||
self._counts[key] = GatewayRequestCounts(
|
||||
successful_requests=existing.successful_requests + counts.successful_requests,
|
||||
failed_requests=existing.failed_requests + counts.failed_requests,
|
||||
)
|
||||
self._counts = dict(fold_counts(chain(self._counts.items(), snapshot.items()))) # mutable-ok: fold replaced
|
||||
|
||||
|
||||
def fold_counts(items: Iterable[tuple[GatewayRequestKey, GatewayRequestCounts]]) -> GatewayRequestSnapshot:
|
||||
"""Sum counts key-wise; the result stays bounded by (date x category x route)."""
|
||||
folded: Final[dict[GatewayRequestKey, GatewayRequestCounts]] = {} # mutable-ok: local fold returned once
|
||||
for key, counts in items:
|
||||
existing = folded.get(key, _EMPTY)
|
||||
folded[key] = GatewayRequestCounts(
|
||||
successful_requests=existing.successful_requests + counts.successful_requests,
|
||||
failed_requests=existing.failed_requests + counts.failed_requests,
|
||||
)
|
||||
return folded
|
||||
|
||||
|
||||
def build_gateway_requests_upsert(snapshot: GatewayRequestSnapshot) -> tuple[str, tuple[str | int, ...]]:
|
||||
"""
|
||||
One ``INSERT ... ON CONFLICT DO UPDATE`` that increments every (date, category,
|
||||
route) in the snapshot. Rows are ordered by the conflict key so concurrent
|
||||
writers lock rows in the same order and cannot deadlock.
|
||||
"""
|
||||
ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route))
|
||||
rows: Final = ", ".join(
|
||||
f"(${base + 1}::text, ${base + 2}::text, ${base + 3}::text, ${base + 4}::bigint, ${base + 5}::bigint, {_UTC_NOW})"
|
||||
for base in range(0, len(ordered) * _COLUMNS_PER_ROW, _COLUMNS_PER_ROW)
|
||||
)
|
||||
sql: Final = (
|
||||
f'INSERT INTO {_TABLE} ("date", "category", "route", "successful_requests", "failed_requests", "updated_at")\n'
|
||||
f"VALUES {rows}\n"
|
||||
'ON CONFLICT ("date", "category", "route") DO UPDATE SET\n'
|
||||
f' "successful_requests" = {_TABLE}."successful_requests" + EXCLUDED."successful_requests",\n'
|
||||
f' "failed_requests" = {_TABLE}."failed_requests" + EXCLUDED."failed_requests",\n'
|
||||
f' "updated_at" = {_UTC_NOW}'
|
||||
)
|
||||
params: Final[tuple[str | int, ...]] = tuple(
|
||||
value
|
||||
for key, counts in ordered
|
||||
for value in (key.date, key.category, key.route, counts.successful_requests, counts.failed_requests)
|
||||
)
|
||||
return sql, params
|
||||
|
||||
|
||||
async def commit_gateway_requests_to_db(
|
||||
|
|
@ -80,50 +138,130 @@ async def commit_gateway_requests_to_db(
|
|||
prisma_client: "PrismaClient",
|
||||
snapshot: GatewayRequestSnapshot,
|
||||
) -> None:
|
||||
"""Upsert one incrementing row per (date, category, route)."""
|
||||
"""Increment every (date, category, route) in the snapshot with a single statement."""
|
||||
if not snapshot:
|
||||
return
|
||||
|
||||
ordered: Final = sorted(snapshot.items(), key=lambda item: (item[0].date, item[0].category, item[0].route))
|
||||
sql, params = build_gateway_requests_upsert(snapshot)
|
||||
await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client
|
||||
|
||||
# pyright: ignore[reportAny] on both lines -- prisma's generated client is untyped,
|
||||
# so .db and every table action off it resolve to Any at this boundary. The dict
|
||||
# literals below are the shape prisma's generated inputs require.
|
||||
async with prisma_client.db.batch_() as batcher: # pyright: ignore[reportAny] # untyped prisma client
|
||||
for key, counts in ordered:
|
||||
columns = asdict(key)
|
||||
batcher.litellm_dailygatewayrequests.upsert( # pyright: ignore[reportAny] # untyped prisma client
|
||||
where={"date_category_route": columns}, # mutable-ok: prisma input is dict-shaped
|
||||
data={ # mutable-ok: prisma input is dict-shaped
|
||||
"create": { # mutable-ok: prisma input is dict-shaped
|
||||
**columns,
|
||||
"successful_requests": counts.successful_requests,
|
||||
"failed_requests": counts.failed_requests,
|
||||
},
|
||||
"update": { # mutable-ok: prisma input is dict-shaped
|
||||
"successful_requests": {"increment": counts.successful_requests}, # mutable-ok: as above
|
||||
"failed_requests": {"increment": counts.failed_requests}, # mutable-ok: as above
|
||||
},
|
||||
},
|
||||
verbose_proxy_logger.debug(
|
||||
"Gateway request tracking - committed %d aggregated rows in one statement", len(snapshot)
|
||||
)
|
||||
|
||||
|
||||
class GatewayRequestRedisBuffer:
|
||||
"""
|
||||
Folds every worker's snapshot through one Redis list so a single pod per
|
||||
interval writes the table, mirroring the spend writer's transaction buffer.
|
||||
|
||||
Each entry is one worker's snapshot as JSON rows; the lock holder pops them,
|
||||
sums them, and commits one statement. A commit failure pushes the summed
|
||||
rows back so the next holder retries, keeping the at-least-once guarantee.
|
||||
If that push fails too, the rows go back to the holder's own accumulator so
|
||||
they ride along with its next flush instead of vanishing with the pop.
|
||||
"""
|
||||
|
||||
def __init__(self, *, redis_cache: RedisCache, pod_lock_manager: PodLockManager) -> None:
|
||||
self._redis_cache: Final = redis_cache
|
||||
self._pod_lock_manager: Final = pod_lock_manager
|
||||
|
||||
async def push(self, snapshot: GatewayRequestSnapshot) -> None:
|
||||
if not snapshot:
|
||||
return
|
||||
rows: Final[_BufferedRows] = tuple(
|
||||
(key.date, key.category, key.route, counts.successful_requests, counts.failed_requests)
|
||||
for key, counts in snapshot.items()
|
||||
)
|
||||
await self._redis_cache.async_rpush(key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, values=(json.dumps(rows),))
|
||||
|
||||
async def _pop_batch(self) -> tuple[str | bytes, ...]:
|
||||
popped: Final[object] = await self._redis_cache.async_lpop( # pyright: ignore[reportAny] # redis returns Any
|
||||
key=REDIS_GATEWAY_REQUESTS_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT
|
||||
)
|
||||
if not popped:
|
||||
return ()
|
||||
return _BUFFERED_ENTRIES.validate_python(popped if isinstance(popped, list) else (popped,))
|
||||
|
||||
async def _pop_all(self) -> AsyncIterator[str | bytes]:
|
||||
while True:
|
||||
batch = await self._pop_batch()
|
||||
for entry in batch:
|
||||
yield entry
|
||||
if len(batch) < MAX_REDIS_BUFFER_DEQUEUE_COUNT:
|
||||
return
|
||||
|
||||
async def pop(self) -> GatewayRequestSnapshot:
|
||||
entries: Final = tuple([entry async for entry in self._pop_all()])
|
||||
return fold_counts(
|
||||
(
|
||||
GatewayRequestKey(date=date, category=category, route=route),
|
||||
GatewayRequestCounts(successful_requests=succeeded, failed_requests=failed),
|
||||
)
|
||||
for entry in entries
|
||||
for date, category, route, succeeded, failed in _BUFFERED_ROWS.validate_json(entry)
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Gateway request tracking - committed %d aggregated rows", len(ordered))
|
||||
async def commit_if_leader(self, prisma_client: "PrismaClient") -> GatewayRequestSnapshot:
|
||||
"""
|
||||
Drain the list and write it as one statement, but only on the pod holding the job lock.
|
||||
|
||||
The lock is a lease, never released: the holder re-enters it on every flush and
|
||||
keeps committing alone until the TTL lapses, so the primary sees one statement
|
||||
per flush interval deployment-wide instead of one per worker.
|
||||
|
||||
Returns the popped rows that could be neither committed nor re-queued, for the
|
||||
caller to keep in memory. Empty on success.
|
||||
"""
|
||||
if not await self._pod_lock_manager.acquire_lock(cronjob_id=GATEWAY_REQUESTS_JOB_NAME):
|
||||
return _NO_COUNTS
|
||||
buffered: Final = await self.pop()
|
||||
try:
|
||||
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=buffered)
|
||||
except Exception: # noqa: BLE001 -- a failed commit must not stop the scheduler
|
||||
verbose_proxy_logger.warning(
|
||||
"Gateway request tracking - failed to commit %d buffered rows, re-queuing to Redis for the next flush",
|
||||
len(buffered),
|
||||
exc_info=True,
|
||||
)
|
||||
return await self._requeue(buffered)
|
||||
return _NO_COUNTS
|
||||
|
||||
async def _requeue(self, snapshot: GatewayRequestSnapshot) -> GatewayRequestSnapshot:
|
||||
try:
|
||||
await self.push(snapshot)
|
||||
except Exception: # noqa: BLE001 -- the rows go back to the caller's accumulator instead
|
||||
verbose_proxy_logger.warning(
|
||||
"Gateway request tracking - Redis re-queue failed, keeping %d rows in memory for the next flush",
|
||||
len(snapshot),
|
||||
exc_info=True,
|
||||
)
|
||||
return snapshot
|
||||
return _NO_COUNTS
|
||||
|
||||
|
||||
async def flush_gateway_requests(
|
||||
prisma_client: "PrismaClient",
|
||||
accumulator: GatewayRequestAccumulator,
|
||||
redis_buffer: GatewayRequestRedisBuffer | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Scheduler entrypoint. Never raises: a metering failure must not kill the job.
|
||||
|
||||
With ``redis_buffer`` the snapshot goes to Redis and only the lease holder
|
||||
writes to Postgres. Shutdown passes no buffer so a departing worker writes its
|
||||
own counts directly instead of parking them behind a lease it may not hold.
|
||||
|
||||
``CancelledError`` is deliberately not caught, so a flush cancelled during
|
||||
shutdown drops its snapshot rather than restoring counts onto an accumulator
|
||||
the process is about to discard.
|
||||
"""
|
||||
snapshot: Final = accumulator.drain()
|
||||
try:
|
||||
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot)
|
||||
if redis_buffer is None:
|
||||
await commit_gateway_requests_to_db(prisma_client=prisma_client, snapshot=snapshot)
|
||||
else:
|
||||
await redis_buffer.push(snapshot)
|
||||
except Exception: # noqa: BLE001 -- a failed flush must not stop the scheduler
|
||||
accumulator.restore(snapshot)
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -131,3 +269,13 @@ async def flush_gateway_requests(
|
|||
len(snapshot),
|
||||
exc_info=True,
|
||||
)
|
||||
return
|
||||
if redis_buffer is None:
|
||||
return
|
||||
try:
|
||||
accumulator.restore(await redis_buffer.commit_if_leader(prisma_client))
|
||||
except Exception: # noqa: BLE001 -- entries still in Redis are drained by the next flush
|
||||
verbose_proxy_logger.warning(
|
||||
"Gateway request tracking - leader drain failed, buffered rows stay in Redis for the next flush",
|
||||
exc_info=True,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -424,6 +424,7 @@ from litellm.proxy.db.exception_handler import (
|
|||
)
|
||||
from litellm.proxy.db.gateway_request_tracking import (
|
||||
GatewayRequestAccumulator,
|
||||
GatewayRequestRedisBuffer,
|
||||
flush_gateway_requests,
|
||||
)
|
||||
from litellm.proxy.db.proxy_worker_heartbeat import (
|
||||
|
|
@ -2355,6 +2356,17 @@ open_telemetry_logger: OpenTelemetry | None = None
|
|||
gateway_request_accumulator: Final = GatewayRequestAccumulator()
|
||||
### INITIALIZE GLOBAL LOGGING OBJECT ###
|
||||
proxy_logging_obj: ProxyLogging = ProxyLogging(user_api_key_cache=user_api_key_cache, premium_user=premium_user)
|
||||
|
||||
|
||||
def _gateway_request_redis_buffer() -> GatewayRequestRedisBuffer | None:
|
||||
"""Shares the spend writer's transaction-buffer Redis and pod lock when use_redis_transaction_buffer is on."""
|
||||
writer: Final = proxy_logging_obj.db_spend_update_writer
|
||||
redis_cache: Final = writer.redis_update_buffer.redis_cache
|
||||
if redis_cache is None or not writer.redis_update_buffer._should_commit_spend_updates_to_redis():
|
||||
return None
|
||||
return GatewayRequestRedisBuffer(redis_cache=redis_cache, pod_lock_manager=writer.pod_lock_manager)
|
||||
|
||||
|
||||
### REDIS QUEUE ###
|
||||
async_result: Final = None
|
||||
celery_app_conn: Final = None
|
||||
|
|
@ -9633,7 +9645,7 @@ class ProxyStartupEvent:
|
|||
flush_gateway_requests,
|
||||
"interval",
|
||||
seconds=batch_writing_interval,
|
||||
args=(prisma_client, gateway_request_accumulator),
|
||||
args=(prisma_client, gateway_request_accumulator, _gateway_request_redis_buffer()),
|
||||
id="update_gateway_requests_job",
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
|
|
|
|||
|
|
@ -8,8 +8,11 @@ from datetime import datetime, timezone
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_GATEWAY_REQUESTS_BUFFER_KEY
|
||||
from litellm.proxy.db.gateway_request_tracking import (
|
||||
GATEWAY_REQUESTS_JOB_NAME,
|
||||
GatewayRequestAccumulator,
|
||||
GatewayRequestRedisBuffer,
|
||||
commit_gateway_requests_to_db,
|
||||
flush_gateway_requests,
|
||||
)
|
||||
|
|
@ -83,40 +86,47 @@ def test_drain_snapshot_is_not_mutated_by_later_records():
|
|||
# ── commit ────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class FakeTable:
|
||||
def __init__(self) -> None:
|
||||
self.upserts: list[dict] = []
|
||||
|
||||
def upsert(self, *, where: dict, data: dict) -> None:
|
||||
self.upserts.append({"where": where, "data": data})
|
||||
|
||||
|
||||
class FakeBatcher:
|
||||
def __init__(self, table: FakeTable) -> None:
|
||||
self.litellm_dailygatewayrequests = table
|
||||
|
||||
async def __aenter__(self) -> "FakeBatcher":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: object) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class FakeDB:
|
||||
def __init__(self, table: FakeTable) -> None:
|
||||
self._table = table
|
||||
def __init__(self) -> None:
|
||||
self.statements: list[tuple[str, tuple[object, ...]]] = []
|
||||
|
||||
def batch_(self) -> FakeBatcher:
|
||||
return FakeBatcher(self._table)
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
self.statements.append((query, args))
|
||||
return len(args) // 5
|
||||
|
||||
|
||||
class FakePrismaClient:
|
||||
def __init__(self) -> None:
|
||||
self.table = FakeTable()
|
||||
self.db = FakeDB(self.table)
|
||||
self.db = FakeDB()
|
||||
|
||||
|
||||
def test_commit_upserts_one_incrementing_row_per_key():
|
||||
def _rows_written(client: FakePrismaClient) -> list[tuple[object, ...]]:
|
||||
"""Every (date, category, route, successful, failed) tuple the database received, in statement order."""
|
||||
return [params[i : i + 5] for _, params in client.db.statements for i in range(0, len(params), 5)]
|
||||
|
||||
|
||||
def test_commit_increments_with_a_single_statement_for_the_whole_snapshot():
|
||||
"""One statement per flush is the whole point: the previous per-key upsert cost
|
||||
the primary (workers x routes) statements per interval."""
|
||||
client = FakePrismaClient()
|
||||
snapshot = {
|
||||
GatewayRequestKey(date="2026-08-01", category="llm", route=route): (
|
||||
GatewayRequestCounts(successful_requests=7, failed_requests=2)
|
||||
)
|
||||
for route in ("/chat/completions", "/embeddings", "/responses", "/v1/messages", "/mcp")
|
||||
}
|
||||
|
||||
asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot))
|
||||
|
||||
assert len(client.db.statements) == 1
|
||||
sql, params = client.db.statements[0]
|
||||
assert sql.count("ON CONFLICT") == 1
|
||||
assert sql.count("(NOW() AT TIME ZONE 'UTC'))") == 5
|
||||
assert len(params) == 25
|
||||
|
||||
|
||||
def test_commit_sql_adds_to_the_existing_row_instead_of_replacing_it():
|
||||
"""A worker only knows its own share; the SQL must add EXCLUDED onto the stored count."""
|
||||
client = FakePrismaClient()
|
||||
snapshot = {
|
||||
GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): (
|
||||
|
|
@ -126,20 +136,36 @@ def test_commit_upserts_one_incrementing_row_per_key():
|
|||
|
||||
asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot))
|
||||
|
||||
assert len(client.table.upserts) == 1
|
||||
written = client.table.upserts[0]
|
||||
assert written["where"] == {
|
||||
"date_category_route": {
|
||||
"date": "2026-08-01",
|
||||
"category": "llm",
|
||||
"route": "/chat/completions",
|
||||
}
|
||||
sql, params = client.db.statements[0]
|
||||
assert 'INSERT INTO "LiteLLM_DailyGatewayRequests"' in sql
|
||||
assert 'ON CONFLICT ("date", "category", "route") DO UPDATE SET' in sql
|
||||
assert (
|
||||
'"successful_requests" = "LiteLLM_DailyGatewayRequests"."successful_requests" + EXCLUDED."successful_requests"'
|
||||
in sql
|
||||
)
|
||||
assert '"failed_requests" = "LiteLLM_DailyGatewayRequests"."failed_requests" + EXCLUDED."failed_requests"' in sql
|
||||
assert params == ("2026-08-01", "llm", "/chat/completions", 7, 2)
|
||||
|
||||
|
||||
def test_commit_placeholders_line_up_with_params():
|
||||
"""$n positions are generated per row; a drift here silently swaps a route for a count."""
|
||||
client = FakePrismaClient()
|
||||
snapshot = {
|
||||
GatewayRequestKey(date="2026-08-01", category="llm", route="/chat/completions"): (
|
||||
GatewayRequestCounts(successful_requests=1, failed_requests=0)
|
||||
),
|
||||
GatewayRequestKey(date="2026-08-01", category="mcp", route="/mcp"): (
|
||||
GatewayRequestCounts(successful_requests=0, failed_requests=3)
|
||||
),
|
||||
}
|
||||
assert written["data"]["update"] == {
|
||||
"successful_requests": {"increment": 7},
|
||||
"failed_requests": {"increment": 2},
|
||||
}
|
||||
assert written["data"]["create"]["successful_requests"] == 7
|
||||
|
||||
asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot))
|
||||
|
||||
sql, params = client.db.statements[0]
|
||||
assert "($1::text, $2::text, $3::text, $4::bigint, $5::bigint," in sql
|
||||
assert "($6::text, $7::text, $8::text, $9::bigint, $10::bigint," in sql
|
||||
assert "$11" not in sql
|
||||
assert params == ("2026-08-01", "llm", "/chat/completions", 1, 0, "2026-08-01", "mcp", "/mcp", 0, 3)
|
||||
|
||||
|
||||
def test_commit_is_deterministically_ordered():
|
||||
|
|
@ -154,17 +180,14 @@ def test_commit_is_deterministically_ordered():
|
|||
|
||||
asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot=snapshot))
|
||||
|
||||
written_order = [
|
||||
(row["where"]["date_category_route"]["date"], row["where"]["date_category_route"]["category"])
|
||||
for row in client.table.upserts
|
||||
]
|
||||
written_order = [(row[0], row[1]) for row in _rows_written(client)]
|
||||
assert written_order == [("2026-08-01", "llm"), ("2026-08-01", "mcp"), ("2026-08-02", "llm")]
|
||||
|
||||
|
||||
def test_commit_skips_the_database_entirely_when_nothing_accumulated():
|
||||
client = FakePrismaClient()
|
||||
asyncio.run(commit_gateway_requests_to_db(prisma_client=client, snapshot={}))
|
||||
assert client.table.upserts == []
|
||||
assert client.db.statements == []
|
||||
|
||||
|
||||
# ── flush ─────────────────────────────────────────────────────────────────────
|
||||
|
|
@ -177,12 +200,12 @@ def test_flush_drains_and_commits():
|
|||
|
||||
asyncio.run(flush_gateway_requests(client, acc))
|
||||
|
||||
assert len(client.table.upserts) == 1
|
||||
assert len(client.db.statements) == 1
|
||||
assert acc.drain() == {}
|
||||
|
||||
|
||||
class ExplodingDB:
|
||||
def batch_(self):
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
raise RuntimeError("db gone")
|
||||
|
||||
|
||||
|
|
@ -208,10 +231,7 @@ def test_failed_flush_keeps_counts_for_the_next_attempt():
|
|||
client = FakePrismaClient()
|
||||
asyncio.run(flush_gateway_requests(client, acc))
|
||||
|
||||
assert client.table.upserts[0]["data"]["update"] == {
|
||||
"successful_requests": {"increment": 1},
|
||||
"failed_requests": {"increment": 1},
|
||||
}
|
||||
assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)]
|
||||
|
||||
|
||||
def test_restored_counts_merge_with_requests_recorded_meanwhile():
|
||||
|
|
@ -223,5 +243,272 @@ def test_restored_counts_merge_with_requests_recorded_meanwhile():
|
|||
client = FakePrismaClient()
|
||||
asyncio.run(flush_gateway_requests(client, acc))
|
||||
|
||||
assert len(client.table.upserts) == 1
|
||||
assert client.table.upserts[0]["data"]["update"]["successful_requests"] == {"increment": 2}
|
||||
assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)]
|
||||
|
||||
|
||||
class ExplodingDBWithInFlightRequest:
|
||||
"""Fails the write after a request has been recorded while it was in flight."""
|
||||
|
||||
def __init__(self, accumulator: GatewayRequestAccumulator) -> None:
|
||||
self.accumulator = accumulator
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
_record(self.accumulator, 500)
|
||||
raise RuntimeError("db gone")
|
||||
|
||||
|
||||
class ExplodingClientWithInFlightRequest:
|
||||
def __init__(self, accumulator: GatewayRequestAccumulator) -> None:
|
||||
self.db = ExplodingDBWithInFlightRequest(accumulator)
|
||||
|
||||
|
||||
def test_restore_keeps_requests_recorded_while_the_failed_write_was_in_flight():
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
asyncio.run(flush_gateway_requests(ExplodingClientWithInFlightRequest(acc), acc))
|
||||
|
||||
client = FakePrismaClient()
|
||||
asyncio.run(flush_gateway_requests(client, acc))
|
||||
|
||||
assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 1, 1)]
|
||||
|
||||
|
||||
# ── redis buffer ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class FakeRedis:
|
||||
def __init__(self) -> None:
|
||||
self.lists: dict[str, list[str]] = {}
|
||||
|
||||
async def async_rpush(self, key: str, values: list[str]) -> int:
|
||||
self.lists.setdefault(key, []).extend(values)
|
||||
return len(self.lists[key])
|
||||
|
||||
async def async_lpop(self, key: str, count: int) -> list[str] | None:
|
||||
queue = self.lists.get(key, [])
|
||||
if not queue:
|
||||
return None
|
||||
popped, self.lists[key] = queue[:count], queue[count:]
|
||||
return popped
|
||||
|
||||
|
||||
class FakePodLock:
|
||||
def __init__(self, *, leader: bool) -> None:
|
||||
self.leader = leader
|
||||
self.held: list[str] = []
|
||||
self.released: list[str] = []
|
||||
|
||||
async def acquire_lock(self, cronjob_id: str) -> bool:
|
||||
self.held.append(cronjob_id)
|
||||
return self.leader
|
||||
|
||||
async def release_lock(self, cronjob_id: str) -> None:
|
||||
self.released.append(cronjob_id)
|
||||
|
||||
|
||||
class FakeLease:
|
||||
"""Redis-side view of the job lock: SET NX by pod id, re-entrant for the holder, freed only by release or TTL."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.holder: str | None = None
|
||||
|
||||
|
||||
class FakeLeasePodLock:
|
||||
def __init__(self, lease: FakeLease, pod_id: str) -> None:
|
||||
self.lease = lease
|
||||
self.pod_id = pod_id
|
||||
|
||||
async def acquire_lock(self, cronjob_id: str) -> bool:
|
||||
if self.lease.holder is None:
|
||||
self.lease.holder = self.pod_id
|
||||
return self.lease.holder == self.pod_id
|
||||
|
||||
async def release_lock(self, cronjob_id: str) -> None:
|
||||
if self.lease.holder == self.pod_id:
|
||||
self.lease.holder = None
|
||||
|
||||
|
||||
def _buffer(redis: FakeRedis, *, leader: bool) -> tuple[GatewayRequestRedisBuffer, FakePodLock]:
|
||||
lock = FakePodLock(leader=leader)
|
||||
return GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=lock), lock # pyright: ignore[reportArgumentType] # duck-typed fakes
|
||||
|
||||
|
||||
def test_non_leader_workers_push_to_redis_and_never_touch_the_database():
|
||||
redis = FakeRedis()
|
||||
client = FakePrismaClient()
|
||||
for _ in range(3):
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
buffer, _ = _buffer(redis, leader=False)
|
||||
asyncio.run(flush_gateway_requests(client, acc, buffer))
|
||||
|
||||
assert client.db.statements == []
|
||||
assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3
|
||||
|
||||
|
||||
def test_leader_folds_every_workers_snapshot_into_one_statement():
|
||||
"""Fifty workers each flushing the same routes must cost the primary one statement, not fifty."""
|
||||
redis = FakeRedis()
|
||||
client = FakePrismaClient()
|
||||
for _ in range(50):
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
_record(acc, 500, route="/responses")
|
||||
buffer, _ = _buffer(redis, leader=False)
|
||||
asyncio.run(flush_gateway_requests(client, acc, buffer))
|
||||
|
||||
leader_acc = GatewayRequestAccumulator()
|
||||
_record(leader_acc, 200)
|
||||
leader, lock = _buffer(redis, leader=True)
|
||||
asyncio.run(flush_gateway_requests(client, leader_acc, leader))
|
||||
|
||||
assert len(client.db.statements) == 1
|
||||
assert _rows_written(client) == [
|
||||
(_today(), "llm", "/chat/completions", 51, 0),
|
||||
(_today(), "llm", "/responses", 0, 50),
|
||||
]
|
||||
assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == []
|
||||
assert lock.held == [GATEWAY_REQUESTS_JOB_NAME]
|
||||
assert lock.released == []
|
||||
|
||||
|
||||
def test_leader_keeps_the_lease_so_staggered_pods_cost_one_statement_per_interval():
|
||||
"""Pods flush on their own clocks; without the lease each one would win the lock in turn and commit alone."""
|
||||
redis = FakeRedis()
|
||||
client = FakePrismaClient()
|
||||
lease = FakeLease()
|
||||
pods = tuple(
|
||||
GatewayRequestRedisBuffer(redis_cache=redis, pod_lock_manager=FakeLeasePodLock(lease, f"pod-{i}")) # pyright: ignore[reportArgumentType] # duck-typed fakes
|
||||
for i in range(4)
|
||||
)
|
||||
|
||||
for _interval in range(3):
|
||||
for pod in pods:
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
asyncio.run(flush_gateway_requests(client, acc, pod))
|
||||
|
||||
assert lease.holder == "pod-0"
|
||||
assert len(client.db.statements) == 3
|
||||
assert [row[3] for row in _rows_written(client)] == [1, 4, 4]
|
||||
assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 3
|
||||
|
||||
|
||||
def test_leader_drains_a_backlog_deeper_than_one_capped_pop():
|
||||
"""More workers than MAX_REDIS_BUFFER_DEQUEUE_COUNT must not leave a growing tail queued behind the cap."""
|
||||
redis = FakeRedis()
|
||||
client = FakePrismaClient()
|
||||
workers = MAX_REDIS_BUFFER_DEQUEUE_COUNT * 2 + 1
|
||||
for _ in range(workers):
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
buffer, _ = _buffer(redis, leader=False)
|
||||
asyncio.run(flush_gateway_requests(client, acc, buffer))
|
||||
|
||||
leader, _ = _buffer(redis, leader=True)
|
||||
asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader))
|
||||
|
||||
assert len(client.db.statements) == 1
|
||||
assert _rows_written(client) == [(_today(), "llm", "/chat/completions", workers, 0)]
|
||||
assert redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == []
|
||||
|
||||
|
||||
def test_leader_with_nothing_buffered_writes_nothing():
|
||||
redis = FakeRedis()
|
||||
client = FakePrismaClient()
|
||||
leader, lock = _buffer(redis, leader=True)
|
||||
|
||||
asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), leader))
|
||||
|
||||
assert client.db.statements == []
|
||||
assert lock.released == []
|
||||
|
||||
|
||||
def test_leader_requeues_to_redis_when_the_database_commit_fails():
|
||||
"""Counts popped from Redis are gone from every worker; a failed commit must put them back."""
|
||||
redis = FakeRedis()
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
_record(acc, 200)
|
||||
leader, lock = _buffer(redis, leader=True)
|
||||
|
||||
asyncio.run(flush_gateway_requests(ExplodingClient(), acc, leader))
|
||||
|
||||
assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1
|
||||
assert lock.released == []
|
||||
assert acc.drain() == {}
|
||||
|
||||
client = FakePrismaClient()
|
||||
retry, _ = _buffer(redis, leader=True)
|
||||
asyncio.run(flush_gateway_requests(client, GatewayRequestAccumulator(), retry))
|
||||
assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)]
|
||||
|
||||
|
||||
class ExplodingRedis(FakeRedis):
|
||||
async def async_rpush(self, key: str, values: list[str]) -> int:
|
||||
raise RuntimeError("redis gone")
|
||||
|
||||
|
||||
class UnreadableRedis(FakeRedis):
|
||||
async def async_lpop(self, key: str, count: int) -> list[str] | None:
|
||||
raise RuntimeError("redis gone mid-flush")
|
||||
|
||||
|
||||
class UnwritableRedis(FakeRedis):
|
||||
"""Pops succeed, pushes fail: a Redis that went read-only between the leader's pop and its re-queue."""
|
||||
|
||||
async def async_rpush(self, key: str, values: list[str]) -> int:
|
||||
raise RuntimeError("redis read-only")
|
||||
|
||||
|
||||
def test_leader_keeps_popped_counts_in_memory_when_both_the_database_and_the_requeue_fail():
|
||||
"""The pop removed the only copy; if Redis will not take it back the leader itself must carry it."""
|
||||
redis = FakeRedis()
|
||||
worker_acc = GatewayRequestAccumulator()
|
||||
_record(worker_acc, 200)
|
||||
_record(worker_acc, 200)
|
||||
worker, _ = _buffer(redis, leader=False)
|
||||
asyncio.run(flush_gateway_requests(FakePrismaClient(), worker_acc, worker))
|
||||
|
||||
degraded = UnwritableRedis()
|
||||
degraded.lists = redis.lists
|
||||
leader_acc = GatewayRequestAccumulator()
|
||||
leader, _ = _buffer(degraded, leader=True)
|
||||
asyncio.run(flush_gateway_requests(ExplodingClient(), leader_acc, leader))
|
||||
assert degraded.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY] == []
|
||||
|
||||
client = FakePrismaClient()
|
||||
retry, _ = _buffer(redis, leader=True)
|
||||
asyncio.run(flush_gateway_requests(client, leader_acc, retry))
|
||||
assert _rows_written(client) == [(_today(), "llm", "/chat/completions", 2, 0)]
|
||||
|
||||
|
||||
def test_leader_whose_redis_read_fails_leaves_the_pushed_rows_for_the_next_flush():
|
||||
"""The scheduler job must not raise, and nothing is popped so nothing needs restoring anywhere."""
|
||||
redis = UnreadableRedis()
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
client = FakePrismaClient()
|
||||
leader, _ = _buffer(redis, leader=True)
|
||||
|
||||
asyncio.run(flush_gateway_requests(client, acc, leader))
|
||||
|
||||
assert client.db.statements == []
|
||||
assert acc.drain() == {}
|
||||
assert len(redis.lists[REDIS_GATEWAY_REQUESTS_BUFFER_KEY]) == 1
|
||||
|
||||
|
||||
def test_failed_redis_push_keeps_counts_locally_for_the_next_flush():
|
||||
acc = GatewayRequestAccumulator()
|
||||
_record(acc, 200)
|
||||
_record(acc, 500)
|
||||
buffer, lock = _buffer(ExplodingRedis(), leader=True)
|
||||
|
||||
asyncio.run(flush_gateway_requests(FakePrismaClient(), acc, buffer))
|
||||
|
||||
assert lock.held == []
|
||||
assert acc.drain() == {
|
||||
GatewayRequestKey(date=_today(), category="llm", route="/chat/completions"): (
|
||||
GatewayRequestCounts(successful_requests=1, failed_requests=1)
|
||||
)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue