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:
mubashir1osmani 2026-07-13 17:26:10 -04:00
parent 6e834e35cb
commit f07a7cbcde
2 changed files with 45 additions and 7 deletions

View file

@ -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

View file

@ -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):