From 3bc7c5f7852ab9a2b3cc585eb688da5e32d97257 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:21:04 -0700 Subject: [PATCH] fix(spend-logs): index error codes with an implied predicate and cap the build size The error code index now covers rows where the error code is not null, which the unchanged filter query implies, so the query stays the same as main and a proxy without the index runs the plan it runs today. The message filter change is reverted. The index build reads every row's metadata, so above about 5M estimated rows the index job skips it with a warning that gives the CREATE INDEX statement, keeping the Helm migration job inside its deadline. An index built by hand under that name is found in place. --- .../request_log_indexes.py | 56 +++++++++++++-- .../spend_management_endpoints.py | 3 - .../migrations/test_request_log_indexes.py | 49 ++++++++++++- .../spend/test_spend_log_error_filters.py | 71 +++++++++++++------ 4 files changed, 149 insertions(+), 30 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 6a261a22d8f..c7d9193a32d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -36,13 +36,17 @@ if TYPE_CHECKING: @dataclass(frozen=True, slots=True) class RequestLogIndex: """One index the migration job owns: the table, the exact Prisma index name, the - column list as it would be written after `ON `, and the WHERE predicate of a - partial index.""" + column list as it would be written after `ON
`, the WHERE predicate of a + partial index, and the most rows the build takes on by itself. An index whose build + has to read each row's metadata costs far more per row than a column index, so above + that estimate it is left for an operator to build off-peak instead of holding the + migration job past its deadline.""" table: str name: str definition: str predicate: str | None = None + max_rows: int | None = None @property def columns(self) -> tuple[str, ...]: @@ -72,7 +76,8 @@ REQUEST_LOG_INDEXES: Final = ( "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_error_code_startTime_idx", "((metadata->'error_information'->>'error_code'), \"startTime\")", - "jsonb_typeof(metadata->'error_information') = 'object'", + "(metadata->'error_information'->>'error_code') IS NOT NULL", + max_rows=5_000_000, ), ) @@ -88,6 +93,10 @@ _CREATE_INDEX_STATEMENT: Final = re.compile( re.IGNORECASE, ) _TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)" +_ESTIMATED_ROWS_SQL: Final = ( + "SELECT COALESCE(SUM(GREATEST(c.reltuples, 0)), 0)::bigint AS rows FROM pg_class c WHERE c.relkind = 'r' " + "AND (c.oid = to_regclass(%(table)s) OR c.oid IN (SELECT relid FROM pg_partition_tree(to_regclass(%(table)s))))" +) _CHILDREN_WITHOUT_THE_INDEX_SQL: Final = ( "SELECT child.relname AS name, n.nspname AS schema, child.relkind = 'p' AS partitioned " "FROM pg_inherits i JOIN pg_class child ON child.oid = i.inhrelid " @@ -143,6 +152,11 @@ class _TableKind: partitioned: bool +@dataclass(frozen=True, slots=True) +class _EstimatedRows: + rows: int + + def filter_request_log_index_diff(diff_sql: str, indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES) -> str: """The `prisma migrate diff` script without the statements that create a migration-job-owned index, which the schema declares and the migrations deliberately do not build.""" @@ -176,8 +190,9 @@ def ensure_request_log_indexes( """Build every listed index that is missing or invalid. Each build step runs under the migration coordinator lock, held per statement so a resolver booting on another replica gets in between partitions rather than waiting for the whole table. Any - failure is logged and left for the next index build; the result says whether - every index ended up valid. Never raises.""" + failure is logged and left for the next index build, and an index over its + `max_rows` is logged and left for an operator; the result says whether every + other index ended up valid. Never raises.""" import psycopg try: @@ -238,11 +253,42 @@ def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: if table is None: logger.info("Table %s does not exist yet, skipping index %s", index.table, index.name) return True + if _too_large_to_build(connection, schema, index): + return True if table.partitioned: return build_index_on_partitioned_table(connection, schema, index) return _build_leaf_index(connection, schema, index.table, index.name, index) +def _too_large_to_build( + connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex +) -> bool: + from psycopg.rows import class_row + + if index.max_rows is None: + return False + existing: Final = _index_state(connection, schema, index.name) + if existing is not None and existing.valid: + return False + with connection.cursor(row_factory=class_row(_EstimatedRows)) as cursor: + rows: Final = sum( + estimate.rows + for estimate in cursor.execute(_ESTIMATED_ROWS_SQL, {"table": _regclass_name(connection, schema, index.table)}) + ) + if rows <= index.max_rows: + return False + logger.warning( + "Not building index %s on %s: it reads about %s rows, more than the %s it builds by itself. Build it off-peak " + "under that name as %s, and the next index build finds it in place", + index.name, + index.table, + rows, + index.max_rows, + index.body, + ) + return True + + def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str: from psycopg import sql diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index d0ec0d072db..3c65fc0bf06 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -96,7 +96,6 @@ _SESSION_REPRESENTATIVE_ORDER_SQL: Final = ( f'call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC, request_id' ) _BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_batch', 'retrieve_batch')" -_HAS_ERROR_INFORMATION_SQL: Final = "jsonb_typeof(metadata->'error_information') = 'object'" _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -2933,8 +2932,6 @@ async def ui_view_spend_logs( sql_conditions.append(f"metadata->>'user_api_key_alias' LIKE ${p}") sql_params.append(f"%{key_alias}%") p += 1 - if error_code is not None or error_message is not None: - sql_conditions.append(_HAS_ERROR_INFORMATION_SQL) if error_code is not None: sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}") sql_params.append(error_code) diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 5ab0e954890..6d28ca0193b 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -8,7 +8,7 @@ import threading import time import uuid from collections.abc import Callable, Iterator, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from types import MappingProxyType from typing import Final @@ -62,6 +62,7 @@ ORIGINAL_MIGRATION_SQL: Final = MappingProxyType( } ) CALL_ID_INDEX_DEFINITION: Final = next(index for index in REQUEST_LOG_INDEXES if index.name == CALL_ID_INDEX) +ERROR_INDEX_DEFINITION: Final = next(index for index in REQUEST_LOG_INDEXES if index.name == ERROR_INDEX) RELEASES: Final = pytest.mark.parametrize( "release", (API_KEY_INDEX_MIGRATION, CALL_ID_INDEX_MIGRATION), ids=("v1.102.1", "v1.103.0") ) @@ -73,6 +74,7 @@ PARTITIONS: Final = MappingProxyType( ) DEFAULT_PARTITION: Final = "LiteLLM_SpendLogs_pdefault" ROWS_PER_PARTITION: Final = 200 +SEEDED_ROWS: Final = 50 RESOLVERS: Final = pytest.mark.parametrize("use_v2_resolver", (True, False), ids=("v2", "v1")) @@ -920,6 +922,51 @@ def test_the_partial_error_index_is_built_beside_the_plain_start_time_index_neve assert _index_predicate(scratch_database, ERROR_INDEX) is not None +def _seed_analyzed_spend_logs(database_url: str, partitioned: bool) -> int: + """Seed the spend logs, partitioned or not, ANALYZE them, and return how many rows they hold.""" + if partitioned: + _partition_spend_logs(database_url) + with psycopg.connect(database_url, autocommit=True) as conn: + for n in range(SEEDED_ROWS): + _insert_spend_log(conn, f"capped-{n}", "2026-10-01") + conn.execute('ANALYZE "LiteLLM_SpendLogs"') + return SEEDED_ROWS + (len(PARTITIONS) + 1) * ROWS_PER_PARTITION if partitioned else SEEDED_ROWS + + +@requires_db +@pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned")) +def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_build( + scratch_database: str, partitioned: bool, caplog: pytest.LogCaptureFixture +) -> None: + """A metadata index costs far more per row to build than a column index, so on a table over its cap the + migration job must finish without it instead of running past its deadline and failing the upgrade.""" + rows: Final = _seed_analyzed_spend_logs(scratch_database, partitioned) + capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=rows - 1) + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True + assert f"Not building index {ERROR_INDEX}" in caplog.text + assert _indexed_table(scratch_database, ERROR_INDEX) is None + + assert ensure_request_log_indexes(scratch_database, "public", (replace(capped, max_rows=rows),)) is True + assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True + + +@requires_db +def test_an_index_an_operator_built_by_hand_on_a_table_over_the_cap_is_found_in_place( + scratch_database: str, caplog: pytest.LogCaptureFixture +) -> None: + _seed_analyzed_spend_logs(scratch_database, partitioned=False) + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute(f'CREATE INDEX CONCURRENTLY "{ERROR_INDEX}" ON "LiteLLM_SpendLogs" {ERROR_INDEX_DEFINITION.body}') + hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX] + + with caplog.at_level("WARNING", logger="litellm_proxy_extras"): + assert ensure_request_log_indexes(scratch_database, "public", (replace(ERROR_INDEX_DEFINITION, max_rows=1),)) + assert "Not building index" not in caplog.text + assert _index_oids(scratch_database)[ERROR_INDEX] == hand_built + + @requires_db def test_an_index_with_another_definition_is_not_taken_for_the_managed_one(scratch_database: str) -> None: with psycopg.connect(scratch_database, autocommit=True) as conn: diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index 04f20bb115f..13bd1405fd1 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -1,8 +1,10 @@ +import time from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import product from pathlib import Path +from types import MappingProxyType from typing import Final import psycopg @@ -13,11 +15,12 @@ from integration._support.process import MIGRATE_DEPLOY, owned_proxy_process from psycopg.rows import class_row from psycopg.types.json import Jsonb -pytestmark: Final = pytest.mark.timeout(240) +pytestmark: Final = pytest.mark.timeout(300) ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_code_startTime_idx" INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" SUCCESS_ROWS: Final = 5000 +STATS_FLUSH_WINDOW_SECONDS: Final = 11.0 @dataclass(frozen=True, slots=True) @@ -39,12 +42,14 @@ ERROR_ROWS: Final = ( ErrorRow("failed-429", "failure", "429", "rate limited by upstream"), ErrorRow("client-disconnected", "success", "499", "Client disconnected the request"), ) -FILTERS: Final = ( - ("error_code", "500", frozenset({"failed-500-overloaded", "failed-500-timeout"})), - ("error_code", "499", frozenset({"client-disconnected"})), - ("error_message", "upstream", frozenset({"failed-500-overloaded", "failed-500-timeout", "failed-429"})), +ERROR_CODES: Final = MappingProxyType( + { + "500": frozenset({"failed-500-overloaded", "failed-500-timeout"}), + "499": frozenset({"client-disconnected"}), + } ) SESSION_VIEWS: Final = ("false", "true") +FILTERS: Final = tuple(product(SESSION_VIEWS, ERROR_CODES)) def _seed(database_url: str, started: datetime) -> None: @@ -92,12 +97,42 @@ def _error_index_scans(database_url: str) -> int: return row.scans +def _settled_error_index_scans(database_url: str, previous: int, unchanged_since: float) -> int: + changed: Final = eventually( + lambda: _error_index_scans(database_url) != previous, + lambda drifted: drifted, + seconds=STATS_FLUSH_WINDOW_SECONDS - (time.monotonic() - unchanged_since), + return_last_on_timeout=True, + ) + if not changed: + return previous + return _settled_error_index_scans(database_url, _error_index_scans(database_url), time.monotonic()) + + def _listed(gateway: Gateway, window: Mapping[str, str], query: Mapping[str, str]) -> frozenset[str]: page: Final = gateway.get("/spend/logs/ui", params={**window, **query}) return frozenset(string_value(object_value(row)["request_id"]) for row in list_value(page["data"])) -def test_error_filters_read_only_the_rows_that_carry_error_information(gateway: Gateway, tmp_path: Path) -> None: +def _list_each_filter( + gateway: Gateway, + database_url: str, + window: Mapping[str, str], + filters: tuple[tuple[str, str], ...], + scans: int, +) -> tuple[tuple[tuple[str, str], frozenset[str], int], ...]: + if not filters: + return () + session, code = filters[0] + listed: Final = _listed(gateway, window, {"error_code": code, "group_by_session": session}) + settled: Final = _settled_error_index_scans(database_url, scans, time.monotonic()) + return ( + ((session, code), listed, settled - scans), + *_list_each_filter(gateway, database_url, window, filters[1:], settled), + ) + + +def test_error_code_filters_read_only_the_rows_that_carry_an_error_code(gateway: Gateway, tmp_path: Path) -> None: now: Final = datetime.now(timezone.utc) window: Final = { "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), @@ -111,19 +146,13 @@ def test_error_filters_read_only_the_rows_that_carry_error_information(gateway: lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=60 ) _seed(database_url, now - timedelta(hours=1)) - scans_before: Final = _error_index_scans(database_url) - - listed: Final = { - (session, name, value): _listed( - booted.gateway, window, {name: value, "group_by_session": session} - ) - for session, (name, value, _) in product(SESSION_VIEWS, FILTERS) - } - assert listed == { - (session, name, value): rows for session, (name, value, rows) in product(SESSION_VIEWS, FILTERS) - } - eventually( - lambda: _error_index_scans(database_url), - lambda scans: scans - scans_before >= len(listed), - seconds=30, + scans: Final = _settled_error_index_scans( + database_url, _error_index_scans(database_url), time.monotonic() ) + + results: Final = _list_each_filter(booted.gateway, database_url, window, FILTERS, scans) + + assert {combo: listed for combo, listed, _ in results} == { + (session, code): ERROR_CODES[code] for session, code in FILTERS + } + assert {combo: index_scans > 0 for combo, _, index_scans in results} == {combo: True for combo in FILTERS}