mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
a6245c1bbb
* test(perf): measure SpendLog-backed UI query latency * test(perf): validate every numeric row and both counter snapshots * test(perf): isolate read API cases and report request metrics --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
110 lines
4.3 KiB
Python
110 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
from collections.abc import Callable, Generator
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Final
|
|
from urllib.parse import urlparse
|
|
|
|
import psycopg
|
|
import pytest
|
|
from psycopg.rows import dict_row
|
|
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
|
|
|
|
|
class ActiveQuery(BaseModel):
|
|
model_config = ConfigDict(frozen=True)
|
|
|
|
pid: int
|
|
leader_pid: int | None
|
|
application_name: str
|
|
state: str | None
|
|
wait_event_type: str | None
|
|
wait_event: str | None
|
|
seconds: float | None
|
|
query: str
|
|
|
|
|
|
def validate_target(url: str) -> None:
|
|
assert os.environ.get("E2E_UI_PERF_DISPOSABLE") == "1", "Disposable database confirmation required"
|
|
assert os.environ.get("E2E_RESET_SPEND_LOGS") != "1", "Never truncate the snapshot baseline"
|
|
assert re.fullmatch(r"db-perf-perf-[a-z0-9-]+\.litellm\.svc", urlparse(url).hostname or ""), (
|
|
"SQL cancellation is restricted to the disposable perf-lane restore"
|
|
)
|
|
|
|
|
|
def active(connection: psycopg.Connection[dict[str, object]]) -> tuple[ActiveQuery, ...]:
|
|
rows: Final = connection.execute("""SELECT pid, leader_pid, application_name, state,
|
|
wait_event_type, wait_event, extract(epoch from clock_timestamp()-query_start)::float AS seconds,
|
|
left(query, 1500) AS query FROM pg_stat_activity
|
|
WHERE datname = current_database() AND pid <> pg_backend_pid()
|
|
AND backend_type IN ('client backend', 'parallel worker')
|
|
AND state IS DISTINCT FROM 'idle' ORDER BY pid""").fetchall()
|
|
return TypeAdapter(tuple[ActiveQuery, ...]).validate_python(rows)
|
|
|
|
|
|
def wait_idle(connection: psycopg.Connection[dict[str, object]], evidence: Path, *, cancel: bool) -> None:
|
|
consecutive = 0
|
|
for queries in (active(connection) for _ in range(30)):
|
|
with evidence.open("a", encoding="utf-8") as output:
|
|
_ = output.write(
|
|
json.dumps(
|
|
{"time": time.time(), "cancel": cancel, "active": tuple(query.model_dump() for query in queries)}
|
|
)
|
|
+ "\n"
|
|
)
|
|
if queries:
|
|
consecutive = 0 # rebind-ok: consecutive idle observations reset on activity
|
|
for query in queries:
|
|
if cancel and query.leader_pid is None:
|
|
_ = connection.execute(
|
|
"SELECT pg_cancel_backend(%s)" if query.state == "active" else "SELECT pg_terminate_backend(%s)",
|
|
(query.pid,),
|
|
)
|
|
else:
|
|
consecutive += 1 # rebind-ok: consecutive idle observations across polling iterations
|
|
if consecutive == 3:
|
|
return
|
|
time.sleep(1)
|
|
pytest.exit("Database did not become idle. No further UI measurements will run", returncode=2)
|
|
|
|
|
|
def settled_counter(read: Callable[[], float], evidence: Path) -> float:
|
|
previous: float | None = None
|
|
consecutive = 0
|
|
for value in (read() for _ in range(10)):
|
|
with evidence.open("a", encoding="utf-8") as output:
|
|
_ = output.write(json.dumps({"time": time.time(), "database_failures": value}) + "\n")
|
|
consecutive = consecutive + 1 if value == previous else 1 # rebind-ok: polling state
|
|
if consecutive == 3:
|
|
return value
|
|
previous = value # rebind-ok: last observed counter for stabilization
|
|
time.sleep(1)
|
|
pytest.exit("Database counters did not settle. No further UI measurements will run", returncode=2)
|
|
|
|
|
|
@contextmanager
|
|
def isolated_database(name: str) -> Generator[None]:
|
|
url: Final = os.environ["DATABASE_URL"]
|
|
validate_target(url)
|
|
evidence: Final = Path(os.environ["E2E_ARTIFACT_DIR"]) / f"ui-isolation-{name}.jsonl"
|
|
with psycopg.Connection[dict[str, object]].connect(
|
|
url,
|
|
autocommit=True,
|
|
connect_timeout=10,
|
|
application_name="ui-isolation-monitor",
|
|
options="-c statement_timeout=5000",
|
|
row_factory=dict_row,
|
|
) as connection:
|
|
try:
|
|
wait_idle(connection, evidence, cancel=False)
|
|
yield
|
|
finally:
|
|
try:
|
|
wait_idle(connection, evidence, cancel=True)
|
|
except Exception as error:
|
|
pytest.exit(f"SQL cleanup failed ({type(error).__name__}). Stopping UI measurements", returncode=2)
|