test(e2e): perf datastore limits count the Postgres statements and Redis commands a request waits on, background work reported on its own (#46798)

* test(e2e): perf datastore limits count what a request waits on, background Postgres and Redis work reported on its own

postgres_calls_per_request and redis_calls_per_request were whole-server totals (pg_stat_database commits, Redis INFO
commands), so spend flushes and config reloads moved them and a new per-request lookup could hide in the noise. They
now count the statements and commands a request waits on before it is forwarded (auth and routing reads, the
rate-limiter script), from pg_stat_statements read before and after the run and a Redis MONITOR connection held
during it. Spend writes, daily aggregates, config reloads and cache invalidations are postgres_background_calls_per_request
and redis_background_calls_per_request. The run output lists the top statements and commands per bucket, and the
counters are read right after the run so the test's own spend sample is not counted

* test(e2e): perf datastore reads only the runtime user's pg_stat_statements rows, other users' rows have a NULL queryid

* test(e2e): perf Redis classifier strips the litellm.caching: namespace the deployed gateway puts on every key

* test(e2e): small_bodies_5rps limits from ten split-metric runs (postgres awaited 1.8, redis awaited 18, postgres background 12.5), spend buffers and cron locks classified as background

* test(e2e): small_bodies_5rps asserts only the awaited Postgres and Redis counts, background lines stay printed

* test(e2e): perf Redis classifier puts the cron lock script's GET and DEL in background, not awaited

---------

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

tests/e2e/perf/cases/small_bodies_5rps.py

@@ -64,11 +64,11 @@ CASE: Final = PerformanceTestCase(
     ),
     limits=(
         Limit("unexpected_status_count", max=3),
-        Limit("gateway_added_seconds", percentile=50, max=0.38),
-        Limit("time_to_first_token_seconds", percentile=99, max=3.0),
+        Limit("gateway_added_seconds", percentile=50, max=0.29),
+        Limit("time_to_first_token_seconds", percentile=99, max=3.7),
         Limit("memory_growth_megabytes", max=100),
-        Limit("postgres_calls_per_request", max=13),
-        Limit("redis_calls_per_request", max=35),
-        Limit("cpu_seconds_per_request", max=0.19),
+        Limit("postgres_calls_per_request", max=1.8),
+        Limit("redis_calls_per_request", max=18),
+        Limit("cpu_seconds_per_request", max=0.29),
     ),
 )

tests/e2e/perf/setup/datastore.py

@@ -1,26 +1,90 @@
-"""Postgres and Redis work counters, read before and after the traffic phase so a case can put a ceiling
-on how much database and cache work the gateway does per request. Both come from the servers' own
-statistics (`pg_stat_database`, Redis `INFO`), so nothing here depends on how the gateway talks to them.
-`DATABASE_URL` and `REDIS_HOST` are the runner's own env, as in the other suites; without them the
-counters are None and any Limit on them reports as missing."""
+"""Postgres and Redis work the gateway did during the traffic phase, split into what a request waits on and
+what runs after the response. Both come from the servers' own accounting, so nothing here depends on how the
+gateway talks to them: Postgres statement counts from `pg_stat_statements` (read before and after the run) and
+Redis commands from a `MONITOR` connection held open for the run. `DATABASE_URL` and `REDIS_HOST` are the
+runner's own env, as in the other suites; without them the counters are None and any Limit on them reports as
+missing. `E2E_DATABASE_ADMIN_URL`, when set, is used once to create the `pg_stat_statements` extension. Only the
+runtime user's own statements are read: other users' rows come back with a NULL queryid and hidden text."""
 
 from __future__ import annotations
 
 import os
+import re
+import threading
+from collections import Counter
+from collections.abc import Generator
+from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Final
+from typing import TYPE_CHECKING, Final, Literal
 
 from pydantic import TypeAdapter
 
-_DB_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()"
-_DB_ROW: Final = TypeAdapter(tuple[int])
+if TYPE_CHECKING:
+    import redis
+
+Bucket = Literal["awaited", "background", "plumbing", "unclassified"]
+
+_COMMITS_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()"
+_STATEMENTS_QUERY: Final = (
+    "SELECT queryid, calls, query FROM pg_stat_statements "
+    "WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database()) "
+    "AND userid = (SELECT usesysid FROM pg_user WHERE usename = current_user) AND queryid IS NOT NULL"
+)
+_CREATE_EXTENSION: Final = "CREATE EXTENSION IF NOT EXISTS pg_stat_statements"
+_COMMITS_ROW: Final = TypeAdapter(tuple[int])
+_STATEMENT_ROWS: Final = TypeAdapter(list[tuple[int, int, str]])
 _INFO: Final = TypeAdapter(dict[str, object])
 _INT: Final = TypeAdapter(int)
+_MONITOR_LINE: Final = TypeAdapter(dict[str, object])
+
+_TABLE: Final = re.compile(r'"(LiteLLM_\w+)"')
+_AWAITED_TABLES: Final = frozenset(
+    {
+        "LiteLLM_VerificationToken",
+        "LiteLLM_UserTable",
+        "LiteLLM_TeamTable",
+        "LiteLLM_TeamMembership",
+        "LiteLLM_OrganizationTable",
+        "LiteLLM_OrganizationMembership",
+        "LiteLLM_EndUserTable",
+        "LiteLLM_BudgetTable",
+        "LiteLLM_ObjectPermissionTable",
+    }
+)
+_PLUMBING_VERBS: Final = frozenset({"BEGIN", "COMMIT", "ROLLBACK", "DEALLOCATE", "SET", "SHOW"})
+_REDIS_READS: Final = frozenset({"GET", "MGET", "HGET", "HGETALL", "EXISTS", "TTL", "PTTL"})
+_REDIS_NAMESPACE: Final = "litellm.caching:"
+_REDIS_WRITES: Final = frozenset(
+    {"SET", "SETEX", "DEL", "EXPIRE", "INCRBYFLOAT", "INCR", "RPUSH", "LPUSH", "LPOP", "RPOP", "LTRIM", "HSET", "SADD"}
+)
+
+
+@dataclass(frozen=True, slots=True)
+class StatementCalls:
+    queryid: int
+    query: str
+    calls: int
+
+    @property
+    def bucket(self) -> Bucket:
+        """A SELECT on the key, user, team, org, end-user or budget tables is what a request waits on at auth
+        and routing. Transaction control and session statements are plumbing. Everything else (spend and
+        audit writes, daily aggregates, config and model reloads) runs after the response or on a timer."""
+        verb: Final = self.query.lstrip("(").split(maxsplit=1)[0].upper() if self.query.strip() else ""
+        tables: Final = _TABLE.findall(self.query)
+        if verb in _PLUMBING_VERBS or not tables:
+            return "plumbing"
+        return "awaited" if verb == "SELECT" and tables[0] in _AWAITED_TABLES else "background"
+
+    @property
+    def summary(self) -> str:
+        return re.sub(r"\s+", " ", self.query)[:110]
 
 
 @dataclass(frozen=True, slots=True)
 class DbCounters:
     commits: int
+    statements: tuple[StatementCalls, ...] | None
 
 
 @dataclass(frozen=True, slots=True)
@@ -29,36 +93,169 @@ class RedisCounters:
     connected_clients: int
 
 
+@dataclass(frozen=True, slots=True)
+class RedisCommand:
+    command: str
+    key: str
+    inside_script: bool
+    count: int
+
+    @property
+    def bucket(self) -> Bucket:
+        """Budget reads, cooldown reads and the rate-limiter script (with every command it issues) happen
+        before the request is forwarded. Spend increments, token counters, config rewrites, cache
+        invalidations (PUBLISH, DEL), the spend buffers the pods flush and any other write happen after the
+        response or on a timer. Reads on keys these rules do not know stay unclassified and count in neither."""
+        key: Final = self.key.removeprefix(_REDIS_NAMESPACE)
+        if key.startswith(("litellm_config:", "daily_", "cronjob_lock:", "litellm_spend", "litellm_daily")):
+            return "background"
+        if self.inside_script or self.command in {"EVALSHA", "EVAL"}:
+            return "awaited"
+        if self.command in {"PUBLISH", "DEL", "SCRIPT", "INFO", "MONITOR"}:
+            return "background"
+        if key.startswith("spend:"):
+            return "awaited" if self.command in _REDIS_READS else "background"
+        if key.startswith("deployment:") or key.endswith(("}:requests", "}:window")):
+            return "awaited"
+        if key.endswith("}:tokens") or self.command in _REDIS_WRITES:
+            return "background"
+        return "unclassified"
+
+    @property
+    def summary(self) -> str:
+        key: Final = re.sub(r"[0-9a-f]{16,}|sk-[A-Za-z0-9_-]+|[0-9a-f-]{36}", "*", self.key)[:70]
+        return f"{self.command} {key}{' (in script)' if self.inside_script else ''}".rstrip()
+
+
+def statement_deltas(before: DbCounters | None, after: DbCounters | None) -> tuple[StatementCalls, ...] | None:
+    """Calls per statement during the run: `after` minus `before` by queryid, new statements whole."""
+    if before is None or after is None or before.statements is None or after.statements is None:
+        return None
+    earlier: Final = {statement.queryid: statement.calls for statement in before.statements}
+    return tuple(
+        StatementCalls(
+            queryid=statement.queryid, query=statement.query, calls=statement.calls - earlier.get(statement.queryid, 0)
+        )
+        for statement in after.statements
+        if statement.calls > earlier.get(statement.queryid, 0)
+    )
+
+
+def calls_in(bucket: Bucket, items: tuple[StatementCalls, ...] | tuple[RedisCommand, ...] | None) -> int | None:
+    if items is None:
+        return None
+    return sum(_count(item) for item in items if item.bucket == bucket)
+
+
+def _count(item: StatementCalls | RedisCommand) -> int:
+    return item.calls if isinstance(item, StatementCalls) else item.count
+
+
 def read_db_counters() -> DbCounters | None:
     url: Final = os.environ.get("DATABASE_URL")
     if url is None:
         return None
     import psycopg
 
+    _ensure_statements_extension(url)
     with psycopg.connect(url) as conn:
-        row: Final = conn.execute(_DB_QUERY).fetchone()
-    (commits,) = _DB_ROW.validate_python(row)
-    return DbCounters(commits=commits)
+        (commits,) = _COMMITS_ROW.validate_python(conn.execute(_COMMITS_QUERY).fetchone())
+        try:
+            rows = _STATEMENT_ROWS.validate_python(conn.execute(_STATEMENTS_QUERY).fetchall())
+        except psycopg.errors.UndefinedTable:
+            return DbCounters(commits=commits, statements=None)
+    return DbCounters(
+        commits=commits,
+        statements=tuple(StatementCalls(queryid=queryid, query=query, calls=calls) for queryid, calls, query in rows),
+    )
 
 
-def read_redis_counters() -> RedisCounters | None:
-    """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are."""
-    host: Final = os.environ.get("REDIS_HOST")
-    if host is None:
-        return None
+def _ensure_statements_extension(url: str) -> None:
+    """Creating the extension needs a superuser, so prefer the admin URL; the runtime user's attempt is
+    allowed to fail (a database without `shared_preload_libraries = pg_stat_statements` cannot have it)."""
+    import psycopg
+
+    try:
+        with psycopg.connect(os.environ.get("E2E_DATABASE_ADMIN_URL", url), autocommit=True) as conn:
+            _ = conn.execute(_CREATE_EXTENSION)
+    except psycopg.Error:
+        return
+
+
+def _redis_client() -> redis.Redis[str]:
     import redis
 
-    client: Final = redis.Redis(
-        host=host,
+    return redis.Redis(
+        host=os.environ["REDIS_HOST"],
         port=int(os.environ.get("REDIS_PORT", "6379")),
         ssl=os.environ.get("REDIS_SSL", "false").lower() == "true",
         password=os.environ.get("REDIS_PASSWORD") or None,
         decode_responses=True,
         socket_connect_timeout=2,
     )
+
+
+def read_redis_counters() -> RedisCounters | None:
+    """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are."""
+    if os.environ.get("REDIS_HOST") is None:
+        return None
+    client: Final = _redis_client()
     stats: Final = _INFO.validate_python(client.info("stats"))
     clients: Final = _INFO.validate_python(client.info("clients"))
     return RedisCounters(
         commands=_INT.validate_python(stats["total_commands_processed"]),
         connected_clients=_INT.validate_python(clients["connected_clients"]),
     )
+
+
+class RedisMonitor:
+    """Every command Redis executed while the monitor was open, counted by command, key and whether a Lua
+    script issued it. `MONITOR` echoes each command to this connection, so it costs Redis one write per
+    command; fine at the rates these cases run."""
+
+    def __init__(self) -> None:
+        self._counts: Counter[tuple[str, str, bool]] = Counter()
+        self._thread: threading.Thread | None = None
+        self._pool: redis.ConnectionPool | None = None
+        self.commands: tuple[RedisCommand, ...] | None = None
+
+    def _run(self) -> None:
+        import redis
+
+        client: Final = _redis_client()
+        self._pool = client.connection_pool  # pyright: ignore[reportAny]  # redis.Redis leaves connection_pool untyped
+        with client.monitor() as monitor:
+            try:
+                while True:
+                    line = _MONITOR_LINE.validate_python(monitor.next_command())
+                    words = str(line["command"]).split(" ", 2)
+                    key = words[1].strip('"') if len(words) > 1 else ""
+                    self._counts[(words[0].upper(), key, line["client_type"] == "lua")] += 1
+            except (redis.ConnectionError, OSError, ValueError):
+                return
+
+    def start(self) -> None:
+        self._thread = threading.Thread(target=self._run, name="redis-monitor", daemon=True)
+        self._thread.start()
+
+    def stop(self) -> tuple[RedisCommand, ...]:
+        if self._pool is not None:
+            self._pool.disconnect()
+        if self._thread is not None:
+            self._thread.join(timeout=5)
+        return tuple(
+            RedisCommand(command=command, key=key, inside_script=inside_script, count=count)
+            for (command, key, inside_script), count in self._counts.items()
+        )
+
+
+@contextmanager
+def redis_monitor() -> Generator[RedisMonitor, None, None]:
+    """Counts Redis commands for the duration of the block; `.commands` is set once the block ends."""
+    monitor: Final = RedisMonitor()
+    if os.environ.get("REDIS_HOST") is not None:
+        monitor.start()
+    try:
+        yield monitor
+    finally:
+        monitor.commands = monitor.stop() if os.environ.get("REDIS_HOST") is not None else None

tests/e2e/perf/setup/observe.py

@@ -15,7 +15,14 @@ from e2e_http import NoBody, get_external, post_external, unwrap
 from memory_readings import RssCapture, RssReading, WorkerKey
 from pydantic import BaseModel
 
-from setup.datastore import DbCounters, RedisCounters
+from setup.datastore import (
+    DbCounters,
+    RedisCommand,
+    RedisCounters,
+    StatementCalls,
+    calls_in,
+    statement_deltas,
+)
 from setup.mocks.reply_common import Provider
 from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase
 from setup.wire import WireRecord
@@ -106,6 +113,7 @@ class Observation:
     db_after: DbCounters | None = None
     redis_before: RedisCounters | None = None
     redis_after: RedisCounters | None = None
+    redis_commands: tuple[RedisCommand, ...] | None = None
     gateway_cpu_seconds: float | None = None
 
     def per_request(self, before: int | None, after: int | None) -> float | None:
@@ -113,6 +121,13 @@ class Observation:
             return None
         return (after - before) / len(self.records)
 
+    def count_per_request(self, count: int | None) -> float | None:
+        return None if count is None or not self.records else count / len(self.records)
+
+    @property
+    def statements(self) -> tuple[StatementCalls, ...] | None:
+        return statement_deltas(self.db_before, self.db_after)
+
     @property
     def retried_markers(self) -> frozenset[str]:
         counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None)
@@ -191,11 +206,13 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio
     if metric == "unexpected_status_count":
         return float(observation.unexpected_status_count())
     if metric == "postgres_calls_per_request":
-        db_before, db_after = observation.db_before, observation.db_after
-        return observation.per_request(db_before and db_before.commits, db_after and db_after.commits)
+        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 == "redis_calls_per_request":
-        redis_before, redis_after = observation.redis_before, observation.redis_after
-        return observation.per_request(redis_before and redis_before.commands, redis_after and redis_after.commands)
+        return observation.count_per_request(calls_in("awaited", observation.redis_commands))
+    if metric == "redis_background_calls_per_request":
+        return observation.count_per_request(calls_in("background", observation.redis_commands))
     if metric == "cpu_seconds_per_request":
         cpu: Final = observation.gateway_cpu_seconds
         return None if cpu is None or not observation.records else cpu / len(observation.records)
@@ -234,7 +251,9 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = (
     ("time_to_first_token_seconds", 99),
     ("memory_growth_megabytes", None),
     ("postgres_calls_per_request", None),
+    ("postgres_background_calls_per_request", None),
     ("redis_calls_per_request", None),
+    ("redis_background_calls_per_request", None),
     ("cpu_seconds_per_request", None),
 )
 
@@ -298,6 +317,7 @@ def render(case: PerformanceTestCase, observation: Observation, rows: list[Row])
         flag = "  OVER" if row.over else "  MISSING" if row.missing else ""
         lines.append(f"  {row.label:32} {observed:>12} {limit:>12}{flag}")
     lines.extend(_diagnostics(observation))
+    lines.extend(_datastore_lines(observation))
     for record in failures:
         lines.append(
             f"  wrong: seq {record.seq} {record.caller} {record.endpoint}{' stream' if record.stream else ''} status {record.status}: {record.error}"
@@ -316,9 +336,49 @@ def _diagnostics(observation: Observation) -> list[str]:
         f"    guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}",
         f"    retried share {len(observation.retried_markers) / count:.4f}",
         f"    spend logged share {'n/a' if share is None else f'{share:.3f}'}",
+        f"    postgres transactions per request {_fmt(observation.per_request(_commits(observation.db_before), _commits(observation.db_after)))}",
+        f"    redis commands per request {_fmt(observation.per_request(_commands(observation.redis_before), _commands(observation.redis_after)))}",
         f"    redis connected clients {_clients(observation.redis_before)} -> {_clients(observation.redis_after)}",
     ]
 
 
+def _datastore_lines(observation: Observation) -> list[str]:
+    """What the per-request Postgres and Redis counts are made of, so a red run names the statement or
+    command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit."""
+    count: Final = max(1, len(observation.records))
+    lines: list[str] = []
+    for title, items in (
+        ("postgres statements", observation.statements),
+        ("redis commands", observation.redis_commands),
+    ):
+        if items is None:
+            lines.append(f"  {title}: n/a")
+            continue
+        for bucket in ("awaited", "background", "plumbing", "unclassified"):
+            in_bucket = sorted((item for item in items if item.bucket == bucket), key=_item_count, reverse=True)
+            total = sum(_item_count(item) for item in in_bucket)
+            if not in_bucket:
+                continue
+            lines.append(f"  {title} {bucket}: {total / count:.3f} per request, {len(in_bucket)} distinct")
+            lines.extend(f"    {_item_count(item) / count:8.3f}  {item.summary}" for item in in_bucket[:8])
+    return lines
+
+
+def _item_count(item: StatementCalls | RedisCommand) -> int:
+    return item.calls if isinstance(item, StatementCalls) else item.count
+
+
+def _commits(counters: DbCounters | None) -> int | None:
+    return None if counters is None else counters.commits
+
+
+def _commands(counters: RedisCounters | None) -> int | None:
+    return None if counters is None else counters.commands
+
+
+def _fmt(value: float | None) -> str:
+    return "n/a" if value is None else f"{value:.3f}"
+
+
 def _clients(counters: RedisCounters | None) -> str:
     return "n/a" if counters is None else str(counters.connected_clients)

tests/e2e/perf/setup/performance_test_case.py

@@ -22,8 +22,10 @@ Metric = Literal[
     "time_to_first_token_seconds",
     "memory_growth_megabytes",
     "postgres_calls_per_request",
+    "postgres_background_calls_per_request",
     "cpu_seconds_per_request",
     "redis_calls_per_request",
+    "redis_background_calls_per_request",
 ]
 
 

tests/e2e/perf/test_perf_cases.py

@@ -9,11 +9,13 @@ import time
 from typing import Final
 
 import pytest
+from cases import CASES
 from idp import Keycloak
 from lifecycle import ResourceManager
 from memory_readings import read_rss_everywhere
 from models import SpendLogsParams
-from cases import CASES
+from proxy_client import ProxyClient
+from setup.datastore import read_db_counters, read_redis_counters, redis_monitor
 from setup.observe import (
     Observation,
     evaluate,
@@ -26,11 +28,9 @@ from setup.observe import (
 )
 from setup.performance_test_case import PerformanceTestCase
 from setup.plan import build_plan
-from setup.datastore import read_db_counters, read_redis_counters
 from setup.podcpu import gateway_cpu_sampler
 from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url
 from setup.run import run_plan
-from proxy_client import ProxyClient
 
 pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
 _SPEND_SAMPLE: Final = 40
@@ -68,9 +68,11 @@ def test_perf_case(
     db_before: Final = read_db_counters()
     redis_before: Final = read_redis_counters()
     run_started: Final = time.monotonic()
-    with gateway_cpu_sampler() as cpu:
+    with gateway_cpu_sampler() as cpu, redis_monitor() as monitor:
         outcome: Final = run_plan(case, plan, setup)
     run_finished: Final = time.monotonic()
+    db_after: Final = read_db_counters()
+    redis_after: Final = read_redis_counters()
     rss_after: Final = read_rss_everywhere(proxy)
     rss_read: Final = time.monotonic()
     assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}"
@@ -81,8 +83,6 @@ def test_perf_case(
         proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE]
     )
     spend_read: Final = time.monotonic()
-    db_after: Final = read_db_counters()
-    redis_after: Final = read_redis_counters()
     mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url))
     guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else ()
     mocks_read: Final = time.monotonic()
@@ -98,6 +98,7 @@ def test_perf_case(
         db_after=db_after,
         redis_before=redis_before,
         redis_after=redis_after,
+        redis_commands=monitor.commands,
         gateway_cpu_seconds=cpu.seconds,
     )
     rows, exceeded_limits = evaluate(case, observation)