mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
test(perf): remove the database_upgrade_5rps e2e case and its harness tests (#50076)
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
02b8da2f16
commit
0d70b9b9a1
4 changed files with 23 additions and 701 deletions
|
|
@ -1,223 +0,0 @@
|
|||
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" and limit.max is not None
|
||||
)
|
||||
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 upgrade_properties(state: UpgradeState, monitor: UpgradeMonitor) -> tuple[tuple[str, float], ...]:
|
||||
"""The upgrade's own numbers under the `perf_` prefix the dashboard ingests, next to the perf rows."""
|
||||
return (
|
||||
("perf_upgrade_schema_seconds", state.schema_seconds),
|
||||
("perf_upgrade_index_seconds", state.index_seconds),
|
||||
("perf_upgrade_overlap_writes", float(monitor.overlap_writes)),
|
||||
("perf_upgrade_max_lock_waiters", float(monitor.max_lock_waiters)),
|
||||
("perf_upgrade_pending_indexes", float(len(state.pending_indexes))),
|
||||
("perf_upgrade_unmatched_requests", float(len(monitor.pending))),
|
||||
)
|
||||
|
||||
|
||||
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),
|
||||
)
|
||||
|
|
@ -1,102 +0,0 @@
|
|||
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,
|
||||
junit_properties,
|
||||
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, upgrade_properties
|
||||
from setup.upgrade_worker import connect, read_state
|
||||
|
||||
pytestmark = pytest.mark.perf
|
||||
|
||||
|
||||
@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, request: pytest.FixtureRequest
|
||||
) -> 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"
|
||||
)
|
||||
rss_after: Final = read_rss_everywhere(proxy)
|
||||
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)
|
||||
item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped
|
||||
assert isinstance(item, pytest.Item)
|
||||
item.user_properties[:] = [
|
||||
entry for entry in item.user_properties if not entry[0].startswith("perf_")
|
||||
] # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set
|
||||
item.user_properties.extend(junit_properties(rows))
|
||||
item.user_properties.extend(upgrade_properties(state, monitor))
|
||||
_ = 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")
|
||||
monitor.assert_complete(state)
|
||||
assert_finished(connection, state)
|
||||
assert {(r.replica, r.hostname) for r in rss_before.readings} == {
|
||||
(r.replica, r.hostname) for r in rss_after.readings
|
||||
}, "Gateway pods replaced during the upgrade"
|
||||
assert not exceeded, "\n".join(exceeded)
|
||||
|
|
@ -8,7 +8,29 @@ from typing import Final
|
|||
import pytest
|
||||
from setup.run import run_until
|
||||
from setup.wire import WireRecord
|
||||
from test_upgrade import record
|
||||
|
||||
|
||||
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}",
|
||||
)
|
||||
|
||||
|
||||
class StopAfterRecord:
|
||||
|
|
|
|||
|
|
@ -1,375 +0,0 @@
|
|||
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,
|
||||
upgrade_properties,
|
||||
)
|
||||
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,))
|
||||
|
||||
|
||||
def _ceiling(metric: str, percentile: int | None = None) -> float:
|
||||
return next(
|
||||
limit.max
|
||||
for limit in upgrade_case().limits
|
||||
if limit.metric == metric and limit.percentile == percentile and limit.max is not None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"over,expected",
|
||||
[
|
||||
("", ()),
|
||||
("unexpected_status_count", ("unexpected_status_count",)),
|
||||
("gateway_added_seconds", ("gateway_added_seconds percentile 50",)),
|
||||
("time_to_first_token_seconds", ("time_to_first_token_seconds percentile 99",)),
|
||||
("memory_growth_megabytes", ("memory_growth_megabytes",)),
|
||||
],
|
||||
)
|
||||
def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples(
|
||||
over: str, expected: tuple[str, ...]
|
||||
) -> None:
|
||||
"""Each sample sits exactly on the case's ceiling, so only the one nudged past it goes over."""
|
||||
failures: Final = int(_ceiling("unexpected_status_count")) + (over == "unexpected_status_count")
|
||||
overhead: Final = _ceiling("gateway_added_seconds", 50) * (1.01 if over == "gateway_added_seconds" else 1)
|
||||
ttft: Final = _ceiling("time_to_first_token_seconds", 99) * (1.01 if over == "time_to_first_token_seconds" else 1)
|
||||
growth: Final = _ceiling("memory_growth_megabytes") + (over == "memory_growth_megabytes")
|
||||
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)
|
||||
|
||||
|
||||
def test_upgrade_properties_carry_the_upgrade_timings_and_what_the_monitor_saw() -> None:
|
||||
state: Final = UpgradeState(
|
||||
phase="complete",
|
||||
started_at=0,
|
||||
schema_seconds=12.5,
|
||||
index_seconds=240.25,
|
||||
required_indexes=("idx_a", "idx_b"),
|
||||
required_migrations=("20260101_a",),
|
||||
pending_indexes=("idx_b",),
|
||||
baseline_rows={},
|
||||
)
|
||||
monitor: Final = UpgradeMonitor(
|
||||
state=lambda: state,
|
||||
check_metrics=lambda: None,
|
||||
read_activity=lambda: ([], 0),
|
||||
read_persisted=lambda records: frozenset(),
|
||||
start_indexes=lambda: None,
|
||||
)
|
||||
monitor.overlap_writes = 37
|
||||
monitor.max_lock_waiters = 3
|
||||
monitor.pending = (record(0), record(1))
|
||||
|
||||
assert dict(upgrade_properties(state, monitor)) == {
|
||||
"perf_upgrade_schema_seconds": 12.5,
|
||||
"perf_upgrade_index_seconds": 240.25,
|
||||
"perf_upgrade_overlap_writes": 37,
|
||||
"perf_upgrade_max_lock_waiters": 3,
|
||||
"perf_upgrade_pending_indexes": 1,
|
||||
"perf_upgrade_unmatched_requests": 2,
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue