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), )