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}