scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-18 04:44:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 125

test(perf): qualify database upgrades with shared 5 RPS traffic (#46802)

Browse source 
* test(perf): qualify database upgrades with shared 5 RPS traffic

* test(perf): share regular latency and memory limits for upgrades

* fix(perf): escape index format placeholder in upgrade worker

* fix(perf): handle uninitialized multiprocess failure counters

* test(perf): match persisted upgrade requests by call ID

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: kerry <kerry@berri.ai>
This commit is contained in:
Kerry Lu 2026-10-17 11:39:00 -07:00 • committed by GitHub
parent 03e3e71efa
commit b5037f28d1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 959 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
tests/e2e/perf/setup/locustfile.py
View file

@ -19,9 +19,10 @@ from typing import Final, cast
import gevent
from locust import FastHttpUser, task
from locust.contrib.fasthttp import FastResponse, ResponseContextManager
from setup.wire import Endpoint, WireRecord, WireRequest, request_body
from pydantic import TypeAdapter
from setup.wire import Endpoint, PlanProgress, WireRecord, WireRequest, request_body, response_id
_PLAN_ADAPTER: Final = TypeAdapter(list[WireRequest])
_JSON: Final = TypeAdapter[object](object)
_PLAN: Final = _PLAN_ADAPTER.validate_json(Path(os.environ["PERF_PLAN_FILE"]).read_bytes())
@ -33,7 +34,8 @@ _INTERVAL_S: Final = _USERS / _RPS
_JITTER: Final = 0.25
_RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan")
_RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8")
_written = 0
_PROGRESS: Final = PlanProgress(len(_PLAN))
_STOP_FILE: Final = os.environ.get("PERF_STOP_FILE")
BLOCK_REASON: Final = "perf-test-block"
_TOKEN_MAX_AGE_S: Final = 240.0
_tokens: dict[str, tuple[str, float]] = {}
@ -155,10 +157,9 @@ def _judge(item: WireRequest, status: int, body_text: str) -> tuple[bool, str |
def _record(record: WireRecord) -> None:
global _written
_ = _RECORDS.write(record.model_dump_json() + "\n")
_RECORDS.flush()
_written += 1
_PROGRESS.completed += 1
def _authorization(user: FastHttpUser, item: WireRequest) -> str:
@ -203,7 +204,7 @@ class PlanUser(FastHttpUser):
a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the
users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate."""
network_timeout = 120.0
network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120"))
connection_timeout = 30.0
def on_start(self) -> None:
@ -221,9 +222,10 @@ class PlanUser(FastHttpUser):
@task
def send(self) -> None:
item: Final = next(_NEXT, None)
take: Final = _PROGRESS.take(stop=_STOP_FILE is not None and Path(_STOP_FILE).exists())
item: Final = next(_NEXT, None) if take else None
if item is None:
if _written >= len(_PLAN):
if _PROGRESS.drained:
_quit(self)
return
started_at: Final = time.time()
@ -237,6 +239,7 @@ class PlanUser(FastHttpUser):
error: str | None = None
header_s = 0.0
response: FastResponse | None = None
identifier: str | None = None # rebind-ok: the ID arrives in a JSON body or the first SSE event
try:
headers: Final = {"Authorization": _authorization(self, item), "Content-Type": "application/json"}
context: ResponseContextManager
@ -265,6 +268,7 @@ class PlanUser(FastHttpUser):
buffer += line
events, buffer = _events(buffer)
for event in events:
identifier = identifier or response_id(event)
if first_token_s is None and token_in_event(item.endpoint, event):
first_token_s = now
text += delta_of_event(item.endpoint, event)
@ -272,7 +276,9 @@ class PlanUser(FastHttpUser):
body_text = context.text or ""
first_byte_s = time.perf_counter() - started
if status == 200:
text = text_of_json(item.endpoint, _as_dict(_JSON.validate_json(body_text)))
payload: Final = _as_dict(_JSON.validate_json(body_text))
identifier = response_id(payload)
text = text_of_json(item.endpoint, payload)
ok, error = _judge(item, status, body_text)
if ok:
context.success()
@ -302,7 +308,8 @@ class PlanUser(FastHttpUser):
mock_delay_s=float(delay_header) if delay_header else None,
call_id=_header(response, "x-litellm-call-id") if response is not None else None,
error=error,
response_id=identifier,
)
)
if _written >= len(_PLAN):
if _PROGRESS.drained:
_quit(self)

111
tests/e2e/perf/setup/run.py
View file

@ -9,9 +9,11 @@ import shutil
import subprocess
import sys
import tempfile
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from typing import BinaryIO, Final
from e2e_config import PROXY_BASE_URL
from pydantic import TypeAdapter
@ -73,7 +75,59 @@ def _keep_records_artifact(records_path: Path) -> None:
_ = shutil.copyfile(records_path, Path(artifact_dir) / "perf-records.jsonl")
def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup: CaseSetup) -> RunOutcome:
class RecordStream:
def __init__(self, source: BinaryIO) -> None:
self.source: Final = source
self.partial = b""
def read(self) -> tuple[WireRecord, ...]:
chunks: Final = (self.partial + self.source.read()).split(b"\n")
self.partial = chunks[-1]
return tuple(_RECORD_ADAPTER.validate_json(line) for line in chunks[:-1] if line)
def run_until(
command: list[str],
env: dict[str, str],
records_path: Path,
stop_path: Path,
stop_when: Callable[[tuple[WireRecord, ...]], bool],
timeout: int,
) -> subprocess.CompletedProcess[str]:
with records_path.open("rb") as records, records_path.with_suffix(".log").open("w+") as output:
with subprocess.Popen(command, env=env, stdout=output, stderr=subprocess.STDOUT, text=True) as process:
deadline: Final = time.monotonic() + timeout
reader: Final = RecordStream(records)
try:
while process.poll() is None:
if stop_when(reader.read()):
stop_path.touch()
if time.monotonic() >= deadline:
raise subprocess.TimeoutExpired(command, timeout)
try:
_ = process.wait(timeout=1)
except subprocess.TimeoutExpired:
pass
finally:
stop_path.touch()
try:
_ = process.wait(timeout=35)
except subprocess.TimeoutExpired:
process.kill()
_ = process.wait(timeout=10)
_keep_records_artifact(records_path)
_ = output.seek(0)
return subprocess.CompletedProcess(command, process.returncode, output.read(), "")
def run_plan(
case: PerformanceTestCase,
plan: tuple[PlannedRequest, ...],
setup: CaseSetup,
*,
stop_when: Callable[[tuple[WireRecord, ...]], bool] | None = None,
request_timeout: float = 120.0,
) -> RunOutcome:
"""Replay `plan` with Locust in a child process and collect its records. The run-time ceiling
is well past the planned duration so Locust's own clock never cuts a slow run short; the
locustfile quits itself once the last record is written."""
@ -82,6 +136,7 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup:
with tempfile.TemporaryDirectory(prefix="e2e-perf-") as directory:
plan_path: Final = Path(directory) / "plan.json"
records_path: Final = Path(directory) / "records.jsonl"
stop_path: Final = Path(directory) / "stop"
_ = plan_path.write_bytes(_PLAN_ADAPTER.dump_json(wire))
records_path.touch()
env: Final = {
@ -91,33 +146,33 @@ def run_plan(case: PerformanceTestCase, plan: tuple[PlannedRequest, ...], setup:
"PERF_RECORDS_FILE": str(records_path),
"PERF_RPS": str(case.traffic.rps),
"PERF_USERS": str(case.traffic.concurrency),
"PERF_REQUEST_TIMEOUT": str(request_timeout),
**({"PERF_STOP_FILE": str(stop_path)} if stop_when is not None else {}),
}
completed: Final = subprocess.run(
[
sys.executable,
"-P",
"-m",
"locust",
"--headless",
"--only-summary",
"--locustfile",
str(_LOCUSTFILE),
"--host",
PROXY_BASE_URL,
"--users",
str(case.traffic.concurrency),
"--spawn-rate",
str(case.traffic.concurrency),
"--run-time",
f"{ceiling_s}s",
"--exit-code-on-error",
"0",
],
env=env,
capture_output=True,
text=True,
timeout=ceiling_s + 120,
check=False,
command: Final = [
sys.executable,
"-P",
"-m",
"locust",
"--headless",
"--only-summary",
"--locustfile",
str(_LOCUSTFILE),
"--host",
PROXY_BASE_URL,
"--users",
str(case.traffic.concurrency),
"--spawn-rate",
str(case.traffic.concurrency),
"--run-time",
f"{ceiling_s}s",
"--exit-code-on-error",
"0",
]
completed: Final = (
run_until(command, env, records_path, stop_path, stop_when, ceiling_s + 120)
if stop_when is not None
else subprocess.run(command, env=env, capture_output=True, text=True, timeout=ceiling_s + 120, check=False)
)
records: Final = tuple(
_RECORD_ADAPTER.validate_json(line)

209
tests/e2e/perf/setup/upgrade.py Normal file
View file

@ -0,0 +1,209 @@
from __future__ import annotations
import math
import time
from collections.abc import Callable
from dataclasses import replace
from typing import Final
import psycopg
from cases.small_bodies_5rps import CASE
from e2e_config import MASTER_KEY
from e2e_http import send_text_external
from prometheus_client.parser import text_string_to_metric_families
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from setup.performance_test_case import PerformanceTestCase
from setup.upgrade_worker import UpgradeState, fingerprints, missing_indexes, read_state, release_indexes
from setup.wire import WireRecord
MAX_TRAFFIC_SECONDS: Final = 6 * 3600
_ROWS: Final = TypeAdapter(tuple[tuple[str, str | None], ...])
_ACTIVITY: Final = TypeAdapter(tuple[list[int], int])
_COUNT: Final = TypeAdapter(tuple[int])
_SERVICES: Final = ("postgres", "batch_write_to_db")
def upgrade_case() -> PerformanceTestCase:
return replace(
CASE,
name="database_upgrade_5rps",
traffic=replace(CASE.traffic, total_requests=int(CASE.traffic.rps * MAX_TRAFFIC_SECONDS)),
limits=tuple(
limit
for limit in CASE.limits
if limit.metric
in {
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
}
),
)
class MetricsHeaders(BaseModel):
model_config = ConfigDict(frozen=True)
authorization: str = Field(default=f"Bearer {MASTER_KEY}", alias="Authorization")
def postgres_failures(text: str) -> float:
families: Final = tuple(text_string_to_metric_families(text))
names: Final = frozenset(family.name for family in families)
required: Final = frozenset(f"litellm_{service}_failed_requests" for service in _SERVICES)
samples: Final = tuple(sample for family in families for sample in family.samples)
has_database_activity: Final = any(
sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0
for sample in samples
)
assert required <= names or has_database_activity, f"Missing database service counters: {required - names}"
failures: Final = tuple(sample.value for sample in samples if sample.name in {f"{name}_total" for name in required})
assert all(math.isfinite(value) and value >= 0 for value in failures), "Invalid database failure counter"
return sum(failures)
def assert_database_metrics(urls: tuple[str, ...]) -> None:
assert urls, "E2E_DATABASE_METRICS_URLS must name every gateway/backend /metrics endpoint"
for url in urls:
assert_database_metrics_url(url)
def assert_database_metrics_url(url: str) -> None:
result: Final = send_text_external("GET", url, headers=MetricsHeaders(), timeout=10)
assert result.status_code == 200, f"Database metrics unavailable at {url}: {result.status_code}"
assert postgres_failures(result.body) == 0, f"Failed PostgreSQL operations reported at {url}"
def persisted(connection: psycopg.Connection[tuple[object, ...]], records: tuple[WireRecord, ...]) -> frozenset[str]:
ids: Final = tuple(record.call_id for record in records if record.call_id)
if not ids:
return frozenset()
rows: Final = _ROWS.validate_python(
connection.execute(
'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=ANY(%s)', (list(ids),)
).fetchall()
)
return matched_response_ids(records, rows)
def matched_response_ids(records: tuple[WireRecord, ...], rows: tuple[tuple[str, str | None], ...]) -> frozenset[str]:
call_ids: Final = frozenset(call_id for _, call_id in rows if call_id)
return frozenset(record.response_id for record in records if record.response_id and record.call_id in call_ids)
def activity(connection: psycopg.Connection[tuple[object, ...]]) -> tuple[list[int], int]:
return _ACTIVITY.validate_python(
connection.execute("""SELECT
ARRAY(SELECT pid FROM pg_stat_progress_create_index WHERE datname=current_database()),
(SELECT count(*) FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock')
""").fetchone()
)
def assert_finished(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None:
assert state.phase == "complete", f"Upgrade did not finish: {state.phase}"
assert not missing_indexes(connection, state.required_indexes), "Required target indexes are missing or invalid"
(missing,) = _COUNT.validate_python(
connection.execute(
"""
SELECT count(*) FROM unnest(%s::text[]) name WHERE NOT EXISTS (
SELECT 1 FROM "_prisma_migrations" WHERE migration_name=name
AND finished_at IS NOT NULL AND rolled_back_at IS NULL)""",
(list(state.required_migrations),),
).fetchone()
)
assert missing == 0, f"{missing} target migrations were not applied"
(unfinished,) = _COUNT.validate_python(
connection.execute("""
SELECT count(*) FROM "_prisma_migrations" WHERE finished_at IS NULL AND rolled_back_at IS NULL
""").fetchone()
)
assert unfinished == 0, f"{unfinished} unfinished migrations"
assert fingerprints(connection, tuple(state.baseline_rows)) == state.baseline_rows, (
"Baseline rows changed or disappeared"
)
class UpgradeMonitor:
def __init__(
self,
*,
state: Callable[[], UpgradeState],
check_metrics: Callable[[], None],
read_activity: Callable[[], tuple[list[int], int]],
read_persisted: Callable[[tuple[WireRecord, ...]], frozenset[str]],
start_indexes: Callable[[], None],
clock: Callable[[], float] = time.monotonic,
wall_clock: Callable[[], float] = time.time,
) -> None:
self.state: Final = state
self.check_metrics: Final = check_metrics
self.read_activity: Final = read_activity
self.read_persisted: Final = read_persisted
self.start_indexes: Final = start_indexes
self.clock: Final = clock
self.wall_clock: Final = wall_clock
self.started: Final = clock()
self.status_limit: Final = next(limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count")
self.pending: tuple[WireRecord, ...] = ()
self.processed = 0
self.successes = 0
self.persisted_count = 0
self.overlap_writes = 0
self.max_lock_waiters = 0
self.previous_builds: frozenset[int] = frozenset()
self.build_seen_at: float | None = None
self.released = False
self.completed_at: float | None = None
def poll(self, records: tuple[WireRecord, ...]) -> bool:
self.check_metrics()
current: Final = self.state()
assert current.phase != "failed", "Target migration/index worker failed"
assert self.clock() - self.started < MAX_TRAFFIC_SECONDS, "Upgrade exceeded its six-hour operational timeout"
successful: Final = tuple(record for record in records if record.status == 200)
self.processed += len(records)
self.successes += len(successful)
assert self.processed - self.successes <= self.status_limit, "Too many unexpected non-200 responses"
assert all(record.ok and record.text for record in successful), "Failed or empty inference response"
assert all(record.response_id and record.call_id for record in successful), (
"Response missing spend-log identifiers"
)
builds, lock_waiters = self.read_activity()
if not self.previous_builds.intersection(builds):
self.build_seen_at = self.wall_clock() if builds else None
self.max_lock_waiters = max(self.max_lock_waiters, lock_waiters)
self.pending += successful
found: Final = self.read_persisted(self.pending)
self.persisted_count += len(found)
still_building, _ = self.read_activity()
since: Final = self.build_seen_at
if since is not None and frozenset(builds).intersection(still_building):
self.overlap_writes += sum(
1 for record in self.pending if record.response_id in found and record.started_at >= since
)
self.previous_builds = frozenset(builds)
self.pending = tuple(record for record in self.pending if record.response_id not in found)
if self.successes and not self.released:
self.start_indexes()
self.released = True
if current.phase == "complete" and self.completed_at is None:
self.completed_at = self.clock()
return self.completed_at is not None and self.clock() - self.completed_at >= 60 and self.successes >= 300
def assert_complete(self, state: UpgradeState) -> None:
assert self.successes >= 300, f"Only {self.successes} requests completed"
assert not self.pending, f"{len(self.pending)} successful requests have no matching spend log"
assert self.persisted_count == self.successes, "Spend-log count does not match successful requests"
assert not state.pending_indexes or self.overlap_writes > 0, "No persisted requests observed during index work"
def monitor_for(connection: psycopg.Connection[tuple[object, ...]], urls: tuple[str, ...]) -> UpgradeMonitor:
return UpgradeMonitor(
state=lambda: read_state(connection),
check_metrics=lambda: assert_database_metrics(urls),
read_activity=lambda: activity(connection),
read_persisted=lambda records: persisted(connection, records),
start_indexes=lambda: release_indexes(connection),
)

146
tests/e2e/perf/setup/upgrade_worker.py Normal file
View file

@ -0,0 +1,146 @@
from __future__ import annotations
import json
import os
import sys
import time
from pathlib import Path
from typing import Final, Literal
import psycopg
from pydantic import BaseModel, ConfigDict, TypeAdapter
class UpgradeState(BaseModel):
model_config = ConfigDict(frozen=True)
phase: Literal["schema", "schema_ready", "indexes", "complete", "failed"]
started_at: float
schema_seconds: float = 0
index_seconds: float = 0
required_indexes: tuple[str, ...]
required_migrations: tuple[str, ...]
pending_indexes: tuple[str, ...]
baseline_rows: dict[str, str]
_STATE: Final = TypeAdapter(tuple[UpgradeState])
_BOOL: Final = TypeAdapter(tuple[bool])
_STRINGS: Final = TypeAdapter(tuple[tuple[str], ...])
_FINGERPRINTS: Final = TypeAdapter(tuple[tuple[str, str], ...])
_FINGERPRINT: Final = """md5(ROW(request_id, api_key, spend, total_tokens, prompt_tokens,
completion_tokens, model, "startTime", "endTime")::text)"""
def connect() -> psycopg.Connection[tuple[object, ...]]:
return psycopg.connect(
os.environ["DATABASE_URL"],
autocommit=True,
connect_timeout=10,
options="-c statement_timeout=10000 -c timezone=UTC",
)
def fingerprints(connection: psycopg.Connection[tuple[object, ...]], ids: tuple[str, ...]) -> dict[str, str]:
rows: Final = connection.execute(
f'SELECT request_id, {_FINGERPRINT} FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', (list(ids),)
).fetchall()
return dict(_FINGERPRINTS.validate_python(rows))
def missing_indexes(connection: psycopg.Connection[tuple[object, ...]], names: tuple[str, ...]) -> tuple[str, ...]:
rows: Final = connection.execute(
"""SELECT name FROM unnest(%s::text[]) AS name WHERE NOT EXISTS (
SELECT 1 FROM pg_index WHERE indexrelid = to_regclass(format('public.%%I', name))
AND indisvalid AND indisready)""",
(list(names),),
).fetchall()
return tuple(row[0] for row in _STRINGS.validate_python(rows))
def read_state(connection: psycopg.Connection[tuple[object, ...]]) -> UpgradeState:
(state,) = _STATE.validate_python(
connection.execute('SELECT state FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone()
)
return state
def write_state(connection: psycopg.Connection[tuple[object, ...]], state: UpgradeState) -> None:
_ = connection.execute(
"""INSERT INTO "_LiteLLMPerfUpgrade" (id, state) VALUES (1, %s::jsonb)
ON CONFLICT (id) DO UPDATE SET state = EXCLUDED.state""",
(state.model_dump_json(),),
)
def release_indexes(connection: psycopg.Connection[tuple[object, ...]]) -> None:
_ = connection.execute('UPDATE "_LiteLLMPerfUpgrade" SET traffic_ready=true WHERE id=1')
def wait_for_traffic(connection: psycopg.Connection[tuple[object, ...]], deadline: float) -> bool:
while time.monotonic() < deadline:
(ready,) = _BOOL.validate_python(
connection.execute('SELECT traffic_ready FROM "_LiteLLMPerfUpgrade" WHERE id=1').fetchone()
)
if ready:
return True
time.sleep(1)
return False
def main() -> int:
assert os.environ.get("E2E_DATABASE_UPGRADE") == "1", "Only run against a disposable restored baseline"
assert os.environ.get("DATABASE_PERF_CONFIRM_DISPOSABLE") == "1", "Disposable database confirmation required"
from litellm_proxy_extras import request_log_indexes
from litellm_proxy_extras.utils import ProxyExtrasDBManager
with connect() as connection:
_ = connection.execute("""CREATE TABLE "_LiteLLMPerfUpgrade" (
id integer PRIMARY KEY CHECK (id=1), state jsonb NOT NULL, traffic_ready boolean NOT NULL DEFAULT false)""")
sample: Final = connection.execute("""
(SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id LIMIT 32)
UNION (SELECT request_id FROM "LiteLLM_SpendLogs" ORDER BY request_id DESC LIMIT 32)""").fetchall()
ids: Final = tuple(row[0] for row in _STRINGS.validate_python(sample))
assert ids, "Restored baseline has no spend logs"
names: Final = tuple(index.name for index in request_log_indexes.REQUEST_LOG_INDEXES)
migrations: Final = Path(request_log_indexes.__file__).parent / "migrations"
initial: Final = UpgradeState(
phase="schema",
started_at=time.time(),
required_indexes=names,
required_migrations=tuple(sorted(path.parent.name for path in migrations.glob("*/migration.sql"))),
pending_indexes=missing_indexes(connection, names),
baseline_rows=fingerprints(connection, ids),
)
assert initial.required_migrations and initial.required_indexes, "Target schema requirements are missing"
write_state(connection, initial)
try:
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True), (
"Schema migration failed"
)
ready: Final = initial.model_copy(
update={
"phase": "schema_ready",
"schema_seconds": time.time() - initial.started_at,
"pending_indexes": missing_indexes(connection, names),
}
)
write_state(connection, ready)
assert wait_for_traffic(connection, time.monotonic() + 1800), "Traffic did not start within 30 minutes"
write_state(connection, ready.model_copy(update={"phase": "indexes"}))
index_start: Final = time.monotonic()
assert ProxyExtrasDBManager.build_request_log_indexes(), "Request-log index build failed"
complete: Final = ready.model_copy(
update={"phase": "complete", "index_seconds": time.monotonic() - index_start}
)
write_state(connection, complete)
_ = sys.stdout.write(complete.model_dump_json() + "\n")
return 0
except Exception as exc:
write_state(connection, read_state(connection).model_copy(update={"phase": "failed"}))
_ = sys.stderr.write(json.dumps({"error": type(exc).__name__}) + "\n")
return 1
if __name__ == "__main__":
raise SystemExit(main())

30
tests/e2e/perf/setup/wire.py
View file

@ -9,12 +9,13 @@ from collections.abc import Mapping
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, Field
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
from setup.corpus.session import Endpoint as CorpusEndpoint
from setup.corpus.session import request_body as corpus_body
Endpoint = Literal["chat_completions", "responses", "messages"]
_PAYLOAD: Final = TypeAdapter(dict[str, object])
class TokenSource(BaseModel):
@ -70,6 +71,33 @@ class WireRecord(BaseModel):
mock_delay_s: float | None
call_id: str | None
error: str | None
response_id: str | None = None
def response_id(payload: Mapping[str, object]) -> str | None:
nested: Final = payload.get("response") or payload.get("message")
source: Final = _PAYLOAD.validate_python(nested) if isinstance(nested, Mapping) else payload
identifier: Final = source.get("id")
return identifier if isinstance(identifier, str) else None
class PlanProgress:
def __init__(self, total: int) -> None:
self.total: Final = total
self.issued = 0
self.completed = 0
self.stopped = False
def take(self, *, stop: bool) -> bool:
self.stopped = self.stopped or stop
if self.stopped or self.issued >= self.total:
return False
self.issued += 1
return True
@property
def drained(self) -> bool:
return (self.stopped or self.issued == self.total) and self.completed == self.issued
_CORPUS_ENDPOINT: Final[Mapping[Endpoint, CorpusEndpoint]] = MappingProxyType(

85
tests/e2e/perf/test_database_upgrade.py Normal file
View file

@ -0,0 +1,85 @@
from __future__ import annotations
import os
import sys
import time
from pathlib import Path
from typing import Final
import pytest
from idp import Keycloak
from lifecycle import ResourceManager
from memory_readings import read_rss_everywhere
from proxy_client import ProxyClient
from setup.datastore import read_db_counters, read_redis_counters
from setup.observe import Observation, evaluate, read_mock_requests, render, reset_mock, set_mock_latency
from setup.plan import build_plan
from setup.podcpu import gateway_cpu_sampler
from setup.provision import MOCK_PROVIDERS, build_setup, mock_url
from setup.run import run_plan
from setup.upgrade import assert_finished, monitor_for, upgrade_case
from setup.upgrade_worker import connect, read_state
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
@pytest.mark.skipif(os.environ.get("E2E_DATABASE_UPGRADE") != "1", reason="Dedicated disposable snapshot lane only")
def test_database_upgrade(proxy: ProxyClient, idp: Keycloak, resources: ResourceManager) -> None:
case: Final = upgrade_case()
urls: Final = tuple(mock_url(provider) for provider in MOCK_PROVIDERS)
for url in urls:
_ = reset_mock(url)
_ = set_mock_latency(url, case.mock_latency)
setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources)
metrics: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip())
with connect() as connection:
assert read_state(connection).phase == "schema_ready", "Expected the target schema ready, index work waiting"
rss_before: Final = read_rss_everywhere(proxy)
db_before: Final = read_db_counters()
redis_before: Final = read_redis_counters()
monitor: Final = monitor_for(connection, metrics)
with gateway_cpu_sampler() as cpu:
outcome: Final = run_plan(case, build_plan(case), setup, stop_when=monitor.poll, request_timeout=30)
assert outcome.locust_exit == 0, outcome.locust_output
remaining: Final = outcome.records[monitor.processed :]
_ = monitor.poll(remaining)
deadline: Final = time.monotonic() + 120
while monitor.pending and time.monotonic() < deadline:
time.sleep(1)
_ = monitor.poll(())
state: Final = read_state(connection)
_ = sys.stdout.write(
f"successful_requests={monitor.successes} persisted_requests={monitor.persisted_count} "
f"unmatched_requests={len(monitor.pending)}\n"
)
monitor.assert_complete(state)
assert_finished(connection, state)
rss_after: Final = read_rss_everywhere(proxy)
assert {r.worker for r in rss_before.readings} == {r.worker for r in rss_after.readings}, (
"Gateway workers restarted"
)
observed: Final = Observation(
records=outcome.records,
mock_requests=tuple(seen for url in urls for seen in read_mock_requests(url)),
guardrail_calls=(),
rss_before=rss_before,
rss_after=rss_after,
wall_seconds=outcome.wall_seconds,
spend_logged_share=1,
db_before=db_before,
db_after=read_db_counters(),
redis_before=redis_before,
redis_after=read_redis_counters(),
gateway_cpu_seconds=cpu.seconds,
)
rows, exceeded = evaluate(case, observed)
_ = sys.stdout.write(render(case, observed, rows) + "\n")
_ = sys.stdout.write(
f"schema_seconds={state.schema_seconds:.3f} index_seconds={state.index_seconds:.3f} "
f"overlap_writes={monitor.overlap_writes} max_lock_waiters={monitor.max_lock_waiters} "
f"pending_indexes={state.pending_indexes}\n"
)
artifact: Final = os.environ.get("E2E_ARTIFACT_DIR")
if artifact:
_ = (Path(artifact) / "database-upgrade.json").write_text(state.model_dump_json(), encoding="utf-8")
assert not exceeded, "\n".join(exceeded)

66
tests/e2e_harness/perf/test_run.py Normal file
View file

@ -0,0 +1,66 @@
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import Final
import pytest
from setup.run import run_until
from setup.wire import WireRecord
from test_upgrade import record
class StopAfterRecord:
def __init__(self, *, partial: Path, resume: Path, fail: bool = False) -> None:
self.partial: Final = partial
self.resume: Final = resume
self.fail: Final = fail
self.received: tuple[WireRecord, ...] = ()
def __call__(self, records: tuple[WireRecord, ...]) -> bool:
self.received += records
if self.partial.exists():
self.resume.touch()
assert not (self.fail and records), "Database failed"
return bool(records)
@pytest.mark.parametrize("fail", [False, True])
def test_controlled_process_preserves_partial_lines_and_drains_on_stop_or_failure(tmp_path: Path, fail: bool) -> None:
records: Final = tmp_path / "records.jsonl"
stop: Final = tmp_path / "stop"
partial: Final = tmp_path / "partial"
resume: Final = tmp_path / "resume"
records.touch()
encoded: Final = record().model_dump_json() + "\n"
program: Final = f"""
import time
from pathlib import Path
with Path({str(records)!r}).open("w") as f:
text = {encoded!r}
f.write(text[:10])
f.flush()
Path({str(partial)!r}).touch()
while not Path({str(resume)!r}).exists():
time.sleep(0.01)
f.write(text[10:])
f.flush()
while not Path({str(stop)!r}).exists():
time.sleep(0.01)
f.write({record(1).model_dump_json()!r} + "\\n")
f.flush()
"""
callback: Final = StopAfterRecord(partial=partial, resume=resume, fail=fail)
command: Final = [sys.executable, "-c", program]
if fail:
with pytest.raises(AssertionError, match="Database failed"):
_ = run_until(command, dict(os.environ), records, stop, callback, timeout=10)
else:
result: Final = run_until(command, dict(os.environ), records, stop, callback, timeout=10)
assert result.returncode == 0, result.stdout
assert callback.received == (record(),)
assert tuple(WireRecord.model_validate_json(line) for line in records.read_text().splitlines()) == (
record(),
record(1),
)

325
tests/e2e_harness/perf/test_upgrade.py Normal file
View file

@ -0,0 +1,325 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Final
import pytest
from cases.small_bodies_5rps import CASE
from memory_readings import RssCapture, RssReading
from setup.observe import Observation, evaluate
from setup.upgrade import MAX_TRAFFIC_SECONDS, UpgradeMonitor, matched_response_ids, postgres_failures, upgrade_case
from setup.upgrade_worker import UpgradeState
from setup.wire import PlanProgress, WireRecord, response_id
def record(seq: int = 0, *, started_at: float = 100) -> WireRecord:
return WireRecord(
seq=seq,
marker=f"request-{seq}",
caller="key",
endpoint="chat_completions",
stream=False,
blocked=False,
started_at=started_at,
status=200,
header_s=1,
first_byte_s=1,
first_token_s=None,
total_s=1,
text="answer",
ok=True,
mock_delay_s=1,
call_id=f"call-{seq}",
error=None,
response_id=f"response-{seq}",
)
def test_spend_logs_match_call_ids_when_response_ids_are_transformed() -> None:
records: Final = (record(0), record(1), record(2))
rows: Final = (("provider-response-id", "call-0"), ("response-1", "unrelated-call"))
assert matched_response_ids(records, rows) == frozenset({"response-0"})
def state(phase: str = "schema_ready", *, pending: bool = True) -> UpgradeState:
return UpgradeState.model_validate(
{
"phase": phase,
"started_at": 0,
"required_indexes": ["index"],
"required_migrations": ["migration"],
"pending_indexes": ["index"] if pending else [],
"baseline_rows": {"seed-row": "hash"},
}
)
class Probe:
def __init__(self) -> None:
self.now = 0.0
self.current = state()
self.builds: list[int] = []
self.found: frozenset[str] = frozenset()
self.releases = 0
self.metric_checks = 0
self.errors = False
def check_metrics(self) -> None:
self.metric_checks += 1
assert not self.errors, "Failed PostgreSQL operation"
def persisted(self, records: tuple[WireRecord, ...]) -> frozenset[str]:
return frozenset(row.response_id for row in records if row.response_id in self.found and row.response_id)
def release(self) -> None:
self.releases += 1
def monitor(self) -> UpgradeMonitor:
return UpgradeMonitor(
state=lambda: self.current,
check_metrics=self.check_metrics,
read_activity=lambda: (self.builds, 2),
read_persisted=self.persisted,
start_indexes=self.release,
clock=lambda: self.now,
wall_clock=lambda: self.now,
)
def test_stop_drains_issued_requests_without_issuing_the_rest() -> None:
progress: Final = PlanProgress(100)
assert progress.take(stop=False)
assert progress.take(stop=False)
progress.completed = 1
assert not progress.take(stop=True)
assert not progress.drained
assert not progress.take(stop=False)
progress.completed = 2
assert progress.drained
assert progress.issued == 2
def test_finite_plan_still_waits_for_its_last_response() -> None:
progress: Final = PlanProgress(1)
assert progress.take(stop=False)
assert not progress.take(stop=False)
assert not progress.drained
progress.completed = 1
assert progress.drained
@pytest.mark.parametrize(
"payload,expected",
[
({"id": "chat-id"}, "chat-id"),
({"type": "response.created", "response": {"id": "response-id"}}, "response-id"),
({"type": "message_start", "message": {"id": "message-id"}}, "message-id"),
({"type": "response.output_item.added", "item": {"id": "item-id"}}, None),
({"type": "content_block_delta", "delta": {"text": "hello"}}, None),
],
)
def test_response_id_uses_response_not_output_item(payload: Mapping[str, object], expected: str | None) -> None:
assert response_id(payload) == expected
def metrics(*, failure: str = "0", batch: str = "0") -> str:
return f"""# TYPE litellm_postgres_failed_requests_total counter
litellm_postgres_failed_requests_total{{function_name="query",error_class="Timeout"}} {failure}
# TYPE litellm_batch_write_to_db_failed_requests_total counter
litellm_batch_write_to_db_failed_requests_total{{function_name="flush",error_class="DatabaseError"}} {batch}
"""
def test_failed_sql_and_async_batch_writes_are_both_counted() -> None:
assert postgres_failures(metrics(failure="2", batch="3")) == 5
assert postgres_failures(metrics()) == 0
def test_no_error_label_series_is_a_valid_zero() -> None:
assert (
postgres_failures("""# TYPE litellm_postgres_failed_requests_total counter
# TYPE litellm_batch_write_to_db_failed_requests_total counter
""")
== 0
)
@pytest.mark.parametrize("service", [None, "postgres", "batch_write_to_db"])
def test_multiprocess_metrics_omit_failure_families_until_a_failure(service: str | None) -> None:
total: Final = """# TYPE litellm_postgres_total_requests_total counter
litellm_postgres_total_requests_total{postgres="postgres"} 1
"""
failure: Final = (
f"# TYPE litellm_{service}_failed_requests_total counter\n"
f'litellm_{service}_failed_requests_total{{function_name="query",error_class="Timeout"}} 1\n'
if service is not None
else ""
)
assert postgres_failures(total + failure) == (1 if service is not None else 0)
@pytest.mark.parametrize("total", ["0", "-1", "NaN", "+Inf"])
def test_missing_failure_families_require_observed_database_activity(total: str) -> None:
with pytest.raises(AssertionError):
_ = postgres_failures(
"# TYPE litellm_postgres_total_requests_total counter\n"
f'litellm_postgres_total_requests_total{{postgres="postgres"}} {total}\n'
)
@pytest.mark.parametrize("body", ["", "# TYPE unrelated counter\n", metrics(failure="NaN"), metrics(batch="-1")])
def test_missing_or_invalid_counters_do_not_pass_as_zero(body: str) -> None:
with pytest.raises(AssertionError):
_ = postgres_failures(body)
def test_index_completion_stops_after_the_post_upgrade_window() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
assert not monitor.poll((record(),))
assert probe.releases == 1
probe.builds = [17]
probe.now = 101
assert not monitor.poll(())
probe.found = frozenset({"response-0", "response-1"})
probe.now = 103
assert not monitor.poll((record(1, started_at=102),))
assert monitor.overlap_writes == 1
assert monitor.persisted_count == 2
assert not monitor.pending
assert monitor.max_lock_waiters == 2
probe.current = state("complete")
assert not monitor.poll(())
probe.now = 162
assert not monitor.poll(())
probe.now = 163
assert not monitor.poll(())
probe.found = frozenset(f"response-{i}" for i in range(300))
assert monitor.poll(tuple(record(i) for i in range(2, 300)))
assert probe.releases == 1
assert probe.metric_checks == 7
def test_http_200_without_spend_row_cannot_complete() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
records: Final = tuple(record(i) for i in range(300))
_ = monitor.poll(records)
with pytest.raises(AssertionError, match="no matching spend log"):
monitor.assert_complete(state("complete", pending=False))
def test_no_pending_index_work_still_requires_300_persisted_requests() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
records: Final = tuple(record(i) for i in range(300))
probe.found = frozenset(f"response-{i}" for i in range(300))
_ = monitor.poll(records)
monitor.assert_complete(state("complete", pending=False))
assert monitor.successes == monitor.persisted_count == 300
with pytest.raises(AssertionError, match="during index work"):
monitor.assert_complete(state("complete", pending=True))
def test_postgres_failure_fails_even_when_inference_succeeded() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
probe.errors = True
with pytest.raises(AssertionError, match="PostgreSQL"):
_ = monitor.poll((record(),))
def test_failed_migration_and_operational_timeout_are_not_success() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
probe.current = state("failed")
with pytest.raises(AssertionError, match="worker failed"):
_ = monitor.poll(())
probe.current = state()
probe.now = MAX_TRAFFIC_SECONDS
with pytest.raises(AssertionError, match="operational timeout"):
_ = monitor.poll(())
def test_upgrade_case_uses_5rps_with_the_same_request_and_memory_limits() -> None:
case: Final = upgrade_case()
assert case.traffic.rps == 5
assert not {"cpu_seconds_per_request", "postgres_calls_per_request", "redis_calls_per_request"}.intersection(
limit.metric for limit in case.limits
)
assert case.limits == tuple(
limit
for limit in CASE.limits
if limit.metric
in {
"unexpected_status_count",
"gateway_added_seconds",
"time_to_first_token_seconds",
"memory_growth_megabytes",
}
)
def failed_record(seq: int) -> WireRecord:
return record(seq).model_copy(update={"status": 503, "ok": False, "text": "", "response_id": None, "call_id": None})
def test_three_non_200s_do_not_start_indexes_or_require_spend_rows() -> None:
probe: Final = Probe()
monitor: Final = probe.monitor()
assert not monitor.poll(tuple(failed_record(i) for i in range(3)))
assert probe.releases == 0
assert monitor.processed == 3
assert monitor.successes == 0
assert not monitor.pending
probe.found = frozenset(f"response-{i}" for i in range(3, 303))
assert not monitor.poll(tuple(record(i) for i in range(3, 303)))
monitor.assert_complete(state("complete", pending=False))
assert probe.releases == 1
assert monitor.processed == 303
assert monitor.successes == monitor.persisted_count == 300
def test_fourth_non_200_fails_across_poll_batches() -> None:
monitor: Final = Probe().monitor()
assert not monitor.poll(tuple(failed_record(i) for i in range(3)))
with pytest.raises(AssertionError, match="Too many unexpected non-200"):
_ = monitor.poll((failed_record(3),))
@pytest.mark.parametrize(
"failed", [record().model_copy(update={"text": ""}), record().model_copy(update={"ok": False})]
)
def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None:
with pytest.raises(AssertionError, match="Failed or empty inference"):
_ = Probe().monitor().poll((failed,))
@pytest.mark.parametrize(
"failures,overhead,ttft,growth,expected",
[
(3, 0.29, 3.7, 100, ()),
(4, 0.29, 3.7, 100, ("unexpected_status_count",)),
(3, 0.30, 3.7, 100, ("gateway_added_seconds percentile 50",)),
(3, 0.29, 3.8, 100, ("time_to_first_token_seconds percentile 99",)),
(3, 0.29, 3.7, 101, ("memory_growth_megabytes",)),
],
)
def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples(
failures: int, overhead: float, ttft: float, growth: float, expected: tuple[str, ...]
) -> None:
nonstream: Final = tuple(
record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100)
)
stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft})
observed: Final = Observation(
records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))),
mock_requests=(),
guardrail_calls=(),
rss_before=RssCapture((RssReading("gateway", "host", 1, 1000),), ()),
rss_after=RssCapture((RssReading("gateway", "host", 1, 1000 + growth),), ()),
wall_seconds=60,
)
rows, exceeded = evaluate(upgrade_case(), observed)
assert tuple(row.label for row in rows if row.over or row.missing) == expected
assert len(exceeded) == len(expected)