mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(e2e): bound spend-log snapshots to a /spend/logs/v2 window
The rate-limited batch spend test snapshotted unattributed rows via the unpaginated /spend/logs whole-table read, which grows with the environment (58MB on stage) and OOMKilled the e2e runner at its 512Mi limit on every scheduled run. Gateway.spend_logs_window pages /spend/logs/v2 over an explicit date window instead, and SpendLogsParams now rejects a filterless read so the whole-table call cannot come back
This commit is contained in:
parent
2166608eb8
commit
8a5ffab1ca
4 changed files with 122 additions and 6 deletions
|
|
@ -17,6 +17,7 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Callable
|
||||
|
||||
import pytest
|
||||
|
|
@ -49,7 +50,7 @@ from e2e_http import (
|
|||
unwrap,
|
||||
)
|
||||
from lifecycle import ResourceManager
|
||||
from models import KeyGenerateBody, SpendLogRow, SpendLogsParams
|
||||
from models import KeyGenerateBody, SpendLogRow
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
|
|
@ -349,6 +350,10 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
the file-read path fires while the batch itself is not blocked.
|
||||
``resources.key()`` cannot set limits, so the key is minted on the gateway
|
||||
directly and its delete deferred.
|
||||
|
||||
Snapshots read /spend/logs/v2 over a bounded window around the test instead
|
||||
of the unpaginated /spend/logs whole-table read, which grows with the
|
||||
environment and OOMed the e2e runner on stage.
|
||||
"""
|
||||
user_id = f"e2e-batch-rl-{unique_marker()}"
|
||||
key = client.gateway.generate_key(
|
||||
|
|
@ -356,8 +361,13 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
)
|
||||
resources.defer(lambda: client.gateway.delete_key(key))
|
||||
|
||||
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
window_end = window_start + timedelta(hours=2)
|
||||
before = frozenset(
|
||||
row.request_id for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
|
||||
row.request_id
|
||||
for row in unattributed_rows(
|
||||
client.gateway.spend_logs_window(start=window_start, end=window_end)
|
||||
)
|
||||
)
|
||||
|
||||
file = unwrap(
|
||||
|
|
@ -379,7 +389,9 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
|
|||
|
||||
new_orphans = [
|
||||
row
|
||||
for row in unattributed_rows(client.gateway.spend_logs(SpendLogsParams()))
|
||||
for row in unattributed_rows(
|
||||
client.gateway.spend_logs_window(start=window_start, end=window_end)
|
||||
)
|
||||
if row.request_id not in before
|
||||
]
|
||||
assert not new_orphans, (
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import time
|
|||
import warnings
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
|
||||
from e2e_http import (
|
||||
NoBody,
|
||||
|
|
@ -50,6 +51,8 @@ from models import (
|
|||
OcrResponse,
|
||||
SpendLogRow,
|
||||
SpendLogs,
|
||||
SpendLogsPage,
|
||||
SpendLogsPageParams,
|
||||
SpendLogsParams,
|
||||
)
|
||||
from e2e_config import (
|
||||
|
|
@ -255,6 +258,28 @@ class Gateway:
|
|||
case _:
|
||||
return []
|
||||
|
||||
def spend_logs_window(self, *, start: datetime, end: datetime) -> list[SpendLogRow]:
|
||||
def fetch(page: int) -> SpendLogsPage:
|
||||
return unwrap(
|
||||
self.transport.get(
|
||||
"/spend/logs/v2",
|
||||
headers=self.transport.master,
|
||||
params=SpendLogsPageParams(
|
||||
start_date=start.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
end_date=end.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
page=page,
|
||||
page_size=100,
|
||||
),
|
||||
response_type=SpendLogsPage,
|
||||
)
|
||||
)
|
||||
|
||||
first = fetch(1)
|
||||
return [
|
||||
*first.data,
|
||||
*(row for page in range(2, first.total_pages + 1) for row in fetch(page).data),
|
||||
]
|
||||
|
||||
def poll_logs_for_key(
|
||||
self, key: str, *, min_rows: int = 1, predicate: RowsPredicate | None = None
|
||||
) -> list[SpendLogRow]:
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from __future__ import annotations
|
|||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, RootModel
|
||||
from pydantic import BaseModel, ConfigDict, RootModel, model_validator
|
||||
|
||||
# ---------- keys ----------
|
||||
|
||||
|
|
@ -255,6 +255,16 @@ class SpendLogsParams(BaseModel):
|
|||
request_id: str | None = None
|
||||
api_key: str | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def require_filter(self) -> SpendLogsParams:
|
||||
if self.request_id is None and self.api_key is None:
|
||||
raise ValueError(
|
||||
"unfiltered /spend/logs returns the entire spend table and OOMs the "
|
||||
"runner on long-lived environments; filter by request_id or api_key, "
|
||||
"or use Gateway.spend_logs_window for a bounded /spend/logs/v2 read"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class SpendLogsPageParams(BaseModel):
|
||||
"""Query for /spend/logs/v2, which requires an explicit date window and
|
||||
|
|
|
|||
|
|
@ -1,17 +1,23 @@
|
|||
"""Unit coverage for the Gateway model-management surface (create_model /
|
||||
delete_model).
|
||||
delete_model) and the bounded spend read-back (spend_logs_window).
|
||||
|
||||
The batches conftest and several llm_translation tests register deployments at
|
||||
runtime through gateway.create_model; when that method went missing, every batch
|
||||
test errored at fixture setup (AttributeError) before a single request reached
|
||||
the proxy. This pins the surface with a typed fake Transport so a rename or
|
||||
signature drift fails here instead of in a live stage run.
|
||||
|
||||
spend_logs_window exists because the unpaginated /spend/logs whole-table read
|
||||
grew past the e2e runner's memory limit on stage and OOMKilled every run; these
|
||||
tests pin its /spend/logs/v2 pagination and that SpendLogsParams can no longer
|
||||
express the unfiltered read.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from batches.batch_client import BatchClient
|
||||
from e2e_gateway import Gateway
|
||||
|
|
@ -30,6 +36,9 @@ from models import (
|
|||
ModelNewBody,
|
||||
ModelNewResponse,
|
||||
ModelsListResponse,
|
||||
SpendLogsPage,
|
||||
SpendLogsPageParams,
|
||||
SpendLogsParams,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -46,6 +55,8 @@ class _RecordingTransport:
|
|||
servable_after_gets: int = 0
|
||||
models_error: UnknownApiError | None = None
|
||||
model_gets: int = 0
|
||||
spend_total: int = 0
|
||||
spend_gets: list[SpendLogsPageParams] = field(default_factory=list)
|
||||
_created: list[str] = field(default_factory=list)
|
||||
|
||||
def post[R: BaseModel](
|
||||
|
|
@ -91,6 +102,22 @@ class _RecordingTransport:
|
|||
return Success(
|
||||
data=response_type.model_validate({"data": [{"id": name} for name in visible]})
|
||||
)
|
||||
if path == "/spend/logs/v2" and response_type is SpendLogsPage:
|
||||
assert isinstance(params, SpendLogsPageParams)
|
||||
self.spend_gets.append(params)
|
||||
offset = (params.page - 1) * params.page_size
|
||||
count = min(params.page_size, max(self.spend_total - offset, 0))
|
||||
return Success(
|
||||
data=response_type.model_validate(
|
||||
{
|
||||
"data": [{"request_id": f"req-{offset + i}"} for i in range(count)],
|
||||
"total": self.spend_total,
|
||||
"page": params.page,
|
||||
"page_size": params.page_size,
|
||||
"total_pages": (self.spend_total + params.page_size - 1) // params.page_size,
|
||||
}
|
||||
)
|
||||
)
|
||||
raise AssertionError(f"unexpected get: {path}")
|
||||
|
||||
def delete[R: BaseModel](
|
||||
|
|
@ -202,3 +229,45 @@ def test_gateway_delete_model_posts_the_model_id() -> None:
|
|||
assert path == "/model/delete"
|
||||
assert isinstance(body, ModelDeleteBody)
|
||||
assert body.id == "registered-id"
|
||||
|
||||
|
||||
WINDOW_START = datetime(2026, 7, 14, 12, 0, 0, tzinfo=timezone.utc)
|
||||
WINDOW_END = datetime(2026, 7, 14, 14, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def test_gateway_spend_logs_window_pages_through_every_row_in_the_window() -> None:
|
||||
transport = _RecordingTransport(spend_total=250)
|
||||
gateway = Gateway(transport=transport)
|
||||
|
||||
rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END)
|
||||
|
||||
assert len(rows) == 250
|
||||
assert len({row.request_id for row in rows}) == 250
|
||||
assert [params.page for params in transport.spend_gets] == [1, 2, 3]
|
||||
assert all(params.start_date == "2026-07-14 12:00:00" for params in transport.spend_gets)
|
||||
assert all(params.end_date == "2026-07-14 14:00:00" for params in transport.spend_gets)
|
||||
|
||||
|
||||
def test_gateway_spend_logs_window_stops_at_an_exact_page_boundary() -> None:
|
||||
transport = _RecordingTransport(spend_total=200)
|
||||
gateway = Gateway(transport=transport)
|
||||
|
||||
rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END)
|
||||
|
||||
assert len(rows) == 200
|
||||
assert [params.page for params in transport.spend_gets] == [1, 2]
|
||||
|
||||
|
||||
def test_gateway_spend_logs_window_returns_empty_for_an_empty_window() -> None:
|
||||
transport = _RecordingTransport(spend_total=0)
|
||||
gateway = Gateway(transport=transport)
|
||||
|
||||
rows = gateway.spend_logs_window(start=WINDOW_START, end=WINDOW_END)
|
||||
|
||||
assert rows == []
|
||||
assert [params.page for params in transport.spend_gets] == [1]
|
||||
|
||||
|
||||
def test_spend_logs_params_rejects_the_unfiltered_whole_table_read() -> None:
|
||||
with pytest.raises(ValidationError, match="spend_logs_window"):
|
||||
SpendLogsParams()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue