test(perf): cap every awaited Postgres statement at 0.5 calls per request (#46944)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/e2e/perf/cases/mixed_size_requests_2rps.py

@@ -69,6 +69,7 @@ CASE: Final = PerformanceTestCase(
         Limit("memory_growth_megabytes", max=600),
         Limit("cpu_seconds_per_request", max=1.26),
         Limit("postgres_calls_per_request", max=3.5),
+        Limit("postgres_max_statement_per_request", max=0.5),
         Limit("redis_calls_per_request", max=21.0),
     ),
 )

tests/e2e/perf/cases/small_bodies_5rps.py

@@ -68,6 +68,7 @@ CASE: Final = PerformanceTestCase(
         Limit("time_to_first_token_seconds", percentile=99, max=3.7),
         Limit("memory_growth_megabytes", max=100),
         Limit("postgres_calls_per_request", max=1.8),
+        Limit("postgres_max_statement_per_request", max=0.5),
         Limit("redis_calls_per_request", max=18),
         Limit("cpu_seconds_per_request", max=0.29),
     ),

tests/e2e/perf/setup/observe.py

@@ -128,6 +128,13 @@ class Observation:
     def statements(self) -> tuple[StatementCalls, ...] | None:
         return statement_deltas(self.db_before, self.db_after)
 
+    @property
+    def busiest_awaited_statement(self) -> StatementCalls | None:
+        """The awaited statement with the most calls: a lookup added to every request shows up here at
+        1.0 per request while cache refreshes stay fractional, which the awaited total cannot tell apart."""
+        awaited: Final = tuple(item for item in self.statements or () if item.bucket == "awaited")
+        return max(awaited, key=_item_count, default=None)
+
     @property
     def retried_markers(self) -> frozenset[str]:
         counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None)
@@ -209,6 +216,11 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio
         return observation.count_per_request(calls_in("awaited", observation.statements))
     if metric == "postgres_background_calls_per_request":
         return observation.count_per_request(calls_in("background", observation.statements))
+    if metric == "postgres_max_statement_per_request":
+        if observation.statements is None:
+            return None
+        busiest: Final = observation.busiest_awaited_statement
+        return observation.count_per_request(0 if busiest is None else busiest.calls)
     if metric == "redis_calls_per_request":
         return observation.count_per_request(calls_in("awaited", observation.redis_commands))
     if metric == "redis_background_calls_per_request":
@@ -225,6 +237,7 @@ class Row:
     percentile: int | None
     observed: float | None
     limit: float | None
+    detail: str = ""
 
     @property
     def label(self) -> str:
@@ -252,6 +265,7 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = (
     ("memory_growth_megabytes", None),
     ("postgres_calls_per_request", None),
     ("postgres_background_calls_per_request", None),
+    ("postgres_max_statement_per_request", None),
     ("redis_calls_per_request", None),
     ("redis_background_calls_per_request", None),
     ("cpu_seconds_per_request", None),
@@ -275,10 +289,11 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[
                 percentile=rank,
                 observed=observe(metric, rank, observation),
                 limit=limit.max if limit else None,
+                detail=_detail(metric, observation),
             )
         )
     exceeded: Final = [
-        f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}"
+        f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}{f' ({row.detail})' if row.detail else ''}"
         if row.over and row.observed is not None and row.limit is not None
         else f"{row.label}: limit set but nothing observed"
         for row in rows
@@ -287,6 +302,13 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[
     return rows, exceeded
 
 
+def _detail(metric: Metric, observation: Observation) -> str:
+    if metric != "postgres_max_statement_per_request":
+        return ""
+    busiest: Final = observation.busiest_awaited_statement
+    return "" if busiest is None else busiest.summary
+
+
 def _row_properties(row: Row) -> tuple[tuple[str, float], ...]:
     observed: Final = ((row.property_name, row.observed),) if row.observed is not None else ()
     limit: Final = ((f"{row.property_name}_limit", row.limit),) if row.limit is not None else ()

tests/e2e/perf/setup/performance_test_case.py

@@ -23,6 +23,7 @@ Metric = Literal[
     "memory_growth_megabytes",
     "postgres_calls_per_request",
     "postgres_background_calls_per_request",
+    "postgres_max_statement_per_request",
     "cpu_seconds_per_request",
     "redis_calls_per_request",
     "redis_background_calls_per_request",

tests/e2e_harness/perf/setup/test_observe.py

@@ -69,5 +69,59 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un
         "perf_gateway_added_seconds_p99": 0.75,
         "perf_postgres_calls_per_request": 10.0,
         "perf_postgres_background_calls_per_request": 0.0,
+        "perf_postgres_max_statement_per_request": 10.0,
         "perf_redis_calls_per_request_limit": 35.0,
     }
+
+
+def _statement(queryid: int, query: str, calls: int) -> StatementCalls:
+    return StatementCalls(queryid=queryid, query=query, calls=calls)
+
+
+def _observation_with(statements: tuple[StatementCalls, ...]) -> Observation:
+    before: Final = tuple(replace(item, calls=0) for item in statements)
+    return Observation(
+        records=tuple(served(seq, 1.0) for seq in range(1000)),
+        mock_requests=(),
+        guardrail_calls=(),
+        rss_before=NO_RSS,
+        rss_after=NO_RSS,
+        wall_seconds=600.0,
+        db_before=DbCounters(commits=0, statements=before),
+        db_after=DbCounters(commits=0, statements=statements),
+    )
+
+
+MAX_STATEMENT_CASE: Final = replace(CASE, limits=(Limit("postgres_max_statement_per_request", max=0.5),))
+TEAM_REFRESH: Final = 'SELECT * FROM "LiteLLM_TeamMembership" WHERE user_id = $1'
+USER_LOOKUP: Final = 'SELECT * FROM "LiteLLM_UserTable" WHERE user_id = $1'
+
+
+def test_cache_refreshes_whose_sum_exceeds_the_cap_pass_the_per_statement_limit() -> None:
+    statements: Final = (
+        _statement(1, TEAM_REFRESH, calls=357),
+        _statement(2, USER_LOOKUP, calls=245),
+        _statement(3, 'INSERT INTO "LiteLLM_DailyModelUsage" VALUES ($1)', calls=1000),
+    )
+
+    rows, exceeded = evaluate(MAX_STATEMENT_CASE, _observation_with(statements))
+
+    assert exceeded == []
+    assert dict(junit_properties(rows))["perf_postgres_max_statement_per_request"] == 0.357
+    assert dict(junit_properties(rows))["perf_postgres_calls_per_request"] == 0.602
+
+
+def test_a_lookup_on_every_request_fails_the_per_statement_limit_and_names_the_statement() -> None:
+    statements: Final = (_statement(1, TEAM_REFRESH, calls=357), _statement(2, USER_LOOKUP, calls=1000))
+
+    _, exceeded = evaluate(MAX_STATEMENT_CASE, _observation_with(statements))
+
+    assert exceeded == [f"postgres_max_statement_per_request: observed 1.0000 > limit 0.5000 ({USER_LOOKUP})"], exceeded
+
+
+def test_per_statement_limit_without_db_counters_reports_nothing_observed() -> None:
+    observation: Final = replace(_observation_with(()), db_before=None, db_after=None)
+
+    _, exceeded = evaluate(MAX_STATEMENT_CASE, observation)
+
+    assert exceeded == ["postgres_max_statement_per_request: limit set but nothing observed"]