mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(e2e): read spend logs via paginated /spend/logs/v2 to stop OOM'ing the runner
The batches suite's test_rate_limited_batch_create_leaves_no_unattributed_spend_row snapshots every unattributed spend row twice through spend_logs(SpendLogsParams()), which used the deprecated, unpaginated /spend/logs. With no filter that endpoint runs find_all over the whole LiteLLM_SpendLogs table, heavy message/response columns included, so against a shared staging database the response OOM-killed the e2e pod mid-test and the run restarted from preflight every time, always dying at that same test Repoint Gateway.spend_logs at /spend/logs/v2, which is paginated, drops the heavy columns, and caps the row count. It walks every page over a bounded default window and returns the same list[SpendLogRow], so the ~15 poll_logs_for_key / spend_logs callers are untouched. v2 requires an explicit date window and matches the token as stored (a SHA-256 hash) rather than the raw sk- key, so spend_logs fills a bounded window when the caller gives none and hashes an sk- key the way the proxy hashes it; the window and pagination fields are added to SpendLogsParams
This commit is contained in:
parent
6e834e35cb
commit
f07a7cbcde
2 changed files with 45 additions and 7 deletions
|
|
@ -8,10 +8,12 @@ Gateway's key/customer methods for cleanup. Read-backs are eventually consistent
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from e2e_http import (
|
||||
NoBody,
|
||||
|
|
@ -49,7 +51,7 @@ from models import (
|
|||
OcrBody,
|
||||
OcrResponse,
|
||||
SpendLogRow,
|
||||
SpendLogs,
|
||||
SpendLogsPage,
|
||||
SpendLogsParams,
|
||||
)
|
||||
from e2e_config import (
|
||||
|
|
@ -65,6 +67,15 @@ from transport import HttpTransport, SplitTransport, Transport
|
|||
RowsPredicate = Callable[[list[SpendLogRow]], bool]
|
||||
|
||||
|
||||
def _hashed_if_raw_key(api_key: str | None) -> str | None:
|
||||
"""The v2 spend-logs filter matches the token as stored on the row (a
|
||||
SHA-256 hash), not the raw ``sk-`` key. Mirror the proxy's ``hash_token`` so
|
||||
a raw key filters correctly; an already-hashed value passes through."""
|
||||
if api_key is None or not api_key.startswith("sk-"):
|
||||
return api_key
|
||||
return hashlib.sha256(api_key.encode()).hexdigest()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Gateway:
|
||||
transport: Transport
|
||||
|
|
@ -243,17 +254,40 @@ class Gateway:
|
|||
# ---- spend read-back ------------------------------------------------
|
||||
|
||||
def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:
|
||||
"""Every spend row matching ``params``, read by paging through the
|
||||
paginated /spend/logs/v2. This replaces the deprecated, unpaginated
|
||||
/spend/logs, which full-table-scans the entire spend history with the
|
||||
heavy message/response columns included and OOM-kills the caller on a
|
||||
large database. v2 needs an explicit window and matches the hashed token
|
||||
the proxy stores, so a bounded default window is filled when the caller
|
||||
gives none and an ``sk-`` key is hashed the way the proxy hashes it."""
|
||||
now = datetime.now(timezone.utc)
|
||||
fmt = "%Y-%m-%d %H:%M:%S"
|
||||
query = params.model_copy(
|
||||
update={
|
||||
"api_key": _hashed_if_raw_key(params.api_key),
|
||||
"start_date": params.start_date or (now - timedelta(days=1)).strftime(fmt),
|
||||
"end_date": params.end_date or (now + timedelta(days=1)).strftime(fmt),
|
||||
}
|
||||
)
|
||||
first = self._spend_logs_page(query, page=1)
|
||||
if first is None:
|
||||
return []
|
||||
rest = (self._spend_logs_page(query, page=page) for page in range(2, first.total_pages + 1))
|
||||
return [*first.data, *(row for page in rest if page is not None for row in page.data)]
|
||||
|
||||
def _spend_logs_page(self, query: SpendLogsParams, *, page: int) -> SpendLogsPage | None:
|
||||
result = self.transport.get(
|
||||
"/spend/logs",
|
||||
"/spend/logs/v2",
|
||||
headers=self.transport.master,
|
||||
params=params,
|
||||
response_type=SpendLogs,
|
||||
params=query.model_copy(update={"page": page}),
|
||||
response_type=SpendLogsPage,
|
||||
)
|
||||
match result:
|
||||
case Success(data=logs):
|
||||
return logs.root
|
||||
case Success(data=page_data):
|
||||
return page_data
|
||||
case _:
|
||||
return []
|
||||
return None
|
||||
|
||||
def poll_logs_for_key(
|
||||
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
|
||||
|
|
|
|||
|
|
@ -274,6 +274,10 @@ class SpendLogs(RootModel[list[SpendLogRow]]):
|
|||
class SpendLogsParams(BaseModel):
|
||||
request_id: str | None = None
|
||||
api_key: str | None = None
|
||||
start_date: str | None = None
|
||||
end_date: str | None = None
|
||||
page: int = 1
|
||||
page_size: int = 100
|
||||
|
||||
|
||||
class SpendLogsPageParams(BaseModel):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue