fix(spend-logs): cap the error index build by table size and log non-blocking partitioned steps

The index job now skips the error code index when the spend log table and its partitions take more than 4 GiB on disk, since build time follows size, not row count. On a partitioned table the logged statements build each partition concurrently and attach it, and boot-time repair leaves the index and its partition copies to those statements.
This commit is contained in:
mateo-berri 2026-10-28 23:33:07 -07:00
parent 29691d6ddd
commit 2616c71da8
4 changed files with 165 additions and 104 deletions

View file

@ -17,6 +17,7 @@ import re
import time
from collections.abc import Callable
from dataclasses import dataclass
from itertools import chain
from typing import TYPE_CHECKING, Final, Literal
from litellm_proxy_extras._logging import logger
@ -37,16 +38,16 @@ if TYPE_CHECKING:
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 <table>`, 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."""
partial index, and the largest table the build takes on by itself. An index whose
build has to read each row's metadata takes time in proportion to the table's size on
disk, so above that size 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
max_bytes: int | None = None
@property
def columns(self) -> tuple[str, ...]:
@ -77,7 +78,7 @@ REQUEST_LOG_INDEXES: Final = (
"LiteLLM_SpendLogs_error_code_startTime_idx",
"((left(metadata->'error_information'->>'error_code', 128)), \"startTime\")",
"(metadata->'error_information'->>'error_code') IS NOT NULL",
max_rows=1_000_000,
max_bytes=4 * 1024**3,
),
)
@ -93,9 +94,8 @@ _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, COALESCE(BOOL_OR(c.reltuples < 0), false) "
"AS unanalyzed FROM pg_class c WHERE c.relkind = 'r' "
_TABLE_BYTES_SQL: Final = (
"SELECT COALESCE(SUM(pg_table_size(c.oid)), 0)::bigint AS total 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 = (
@ -154,9 +154,8 @@ class _TableKind:
@dataclass(frozen=True, slots=True)
class _EstimatedRows:
rows: int
unanalyzed: bool
class _TableBytes:
total: int
_IndexOutcome = Literal["in_place", "left_for_operator", "not_in_place"]
@ -195,8 +194,8 @@ 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, and an index over its
`max_rows` is logged and left for an operator; the result says whether every
failure is logged and left for the next index build, and an index on a table over its
`max_bytes` is logged and left for an operator; the result says whether every
other index ended up valid. Never raises."""
import psycopg
@ -281,73 +280,94 @@ def _ensure_index(
def _too_large_to_build(
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex, partitioned: bool
) -> bool:
if index.max_rows is None:
if index.max_bytes is None:
return False
existing: Final = _index_state(connection, schema, index.name)
if existing is not None and existing.valid:
return False
rows: Final = _estimated_rows(connection, schema, index.table)
if rows <= index.max_rows:
size: Final = _table_bytes(connection, schema, index.table)
if size <= index.max_bytes:
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 "
"and the next index build finds it in place: %s",
"Not building index %s on %s: the table takes %s on disk, more than the %s it builds by itself. Build it "
"off-peak by running each of these statements on its own, and the next index build finds it in place: %s",
index.name,
index.table,
rows,
index.max_rows,
_operator_build_statement(
connection,
schema,
index,
partitioned,
interrupted=existing is not None and existing.table == index.table,
),
_gib(size),
_gib(index.max_bytes),
"; ".join(_operator_build_statements(connection, schema, index, index.table, index.name, partitioned)),
)
return True
def _estimated_rows(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> int:
"""The planner's row count for the table and its partitions. A table never analyzed counts zero rows however
many it holds (a freshly restored database), so it is analyzed first."""
from psycopg import sql
estimate: Final = _planner_estimate(connection, schema, table)
if not estimate.unanalyzed:
return estimate.rows
connection.execute(sql.SQL("ANALYZE {}").format(sql.Identifier(schema, table)))
return _planner_estimate(connection, schema, table).rows
def _gib(size: int) -> str:
return f"{size / 1024**3:.1f} GiB"
def _planner_estimate(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> _EstimatedRows:
def _table_bytes(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> int:
"""The table's size on disk across its partitions, TOAST and dead rows included: an upper bound on what a build
reads, known without ANALYZE, where a row count misses how large each row's metadata is."""
from psycopg.rows import class_row
with connection.cursor(row_factory=class_row(_EstimatedRows)) as cursor:
estimate: Final = cursor.execute(
_ESTIMATED_ROWS_SQL, {"table": _regclass_name(connection, schema, table)}
).fetchone()
return estimate if estimate is not None else _EstimatedRows(rows=0, unanalyzed=False)
with connection.cursor(row_factory=class_row(_TableBytes)) as cursor:
size: Final = cursor.execute(_TABLE_BYTES_SQL, {"table": _regclass_name(connection, schema, table)}).fetchone()
return size.total if size is not None else 0
def _operator_build_statement(
def _operator_build_statements(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
index: RequestLogIndex,
table: str,
name: str,
partitioned: bool,
interrupted: bool,
) -> str:
"""What an operator runs by hand: CONCURRENTLY wherever Postgres allows it, and a REINDEX of the invalid index an
interrupted build left, which a CREATE ... IF NOT EXISTS would silently skip."""
) -> tuple[str, ...]:
"""What an operator runs by hand to finish the build from where it stands: the steps this job takes, so no
statement holds spend-log writes back for the length of a build. A partitioned table gets the metadata-only
parent index, then each partition still missing its child gets a CONCURRENTLY build and ATTACH PARTITION. An
invalid index an interrupted build left gets REINDEX CONCURRENTLY, which CREATE ... IF NOT EXISTS would skip."""
from psycopg import sql
if interrupted:
reindex: Final = "REINDEX INDEX {}" if partitioned else "REINDEX INDEX CONCURRENTLY {}"
return sql.SQL(reindex).format(sql.Identifier(schema, index.name)).as_string(connection)
build: Final = (
"CREATE INDEX IF NOT EXISTS {} ON {} " if partitioned else "CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} "
found: Final = _index_state(connection, schema, name)
existing: Final = found if found is not None and found.table == table else None
if existing is not None and existing.valid:
return ()
if not partitioned and existing is not None:
return (sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)).as_string(connection),)
if not partitioned:
leaf: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)
)
return (_create_index_statement(connection, leaf, index.body).decode(),)
parent: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)
)
prefix: Final = sql.SQL(build).format(sql.Identifier(index.name), sql.Identifier(schema, index.table))
return _create_index_statement(connection, prefix, index.body).decode()
parent_statement: Final = (
() if existing is not None else (_create_index_statement(connection, parent, index.body).decode(),)
)
children: Final = _children_without_the_index(connection, schema, table, name)
return parent_statement + tuple(
chain.from_iterable(_operator_child_statements(connection, schema, name, child, index) for child in children)
)
def _operator_child_statements(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
parent_index: str,
child: _Relation,
index: RequestLogIndex,
) -> tuple[str, ...]:
from psycopg import sql
child_index: Final = index.partition_index_name(child.name)
attach: Final = sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format(
sql.Identifier(schema, parent_index), sql.Identifier(child.schema, child_index)
)
built: Final = _operator_build_statements(
connection, child.schema, index, child.name, child_index, child.partitioned
)
return (*built, attach.as_string(connection))
def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str:

View file

@ -74,11 +74,10 @@ _MIGRATION_DEADLOCK_MARKER = "deadlock detected"
_MIGRATION_LOCK_TIMEOUT_MARKER = "canceling statement due to lock timeout"
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
_ROW_CAPPED_REQUEST_LOG_INDEXES: Final = frozenset(
index.name for index in REQUEST_LOG_INDEXES if index.max_rows is not None
)
_SIZE_CAPPED_REQUEST_LOG_INDEXES: Final = tuple(index for index in REQUEST_LOG_INDEXES if index.max_bytes is not None)
_INVALID_LITELLM_INDEXES_SQL: Final = (
"SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
"SELECT n.nspname AS schema, c.relname AS name, pg_size_pretty(pg_table_size(t.oid)) AS table_size, "
't.relname AS "table" '
"FROM pg_index i "
"JOIN pg_class c ON c.oid = i.indexrelid "
"JOIN pg_class t ON t.oid = i.indrelid "
@ -97,6 +96,11 @@ class _InvalidIndex:
schema: str
name: str
table_size: str
table: str
@property
def left_to_the_request_log_index_build(self) -> bool:
return any(self.name == index.partition_index_name(self.table) for index in _SIZE_CAPPED_REQUEST_LOG_INDEXES)
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@ -853,8 +857,11 @@ class ProxyExtrasDBManager:
def _invalid_litellm_indexes(
conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
) -> tuple[_InvalidIndex, ...]:
rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
return tuple(_InvalidIndex(*row) for row in rows if row[1] not in _ROW_CAPPED_REQUEST_LOG_INDEXES)
from psycopg.rows import class_row
with conn.cursor(row_factory=class_row(_InvalidIndex)) as cursor:
invalid: Final = cursor.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
return tuple(index for index in invalid if not index.left_to_the_request_log_index_build)
@staticmethod
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:

View file

@ -923,26 +923,31 @@ 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."""
def _seed_spend_logs(database_url: str, partitioned: bool) -> None:
"""Seed the spend logs, partitioned or not, and never ANALYZE them, as in a freshly restored database."""
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
def _run_the_logged_statements(database_url: str, messages: list[str]) -> None:
logged: Final = next(message for message in messages if "Not building index" in message)
with psycopg.connect(database_url, autocommit=True) as conn:
for statement in logged.split("finds it in place: ", 1)[1].split("; "):
conn.execute(statement.encode())
@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(
def test_an_index_on_a_table_over_its_size_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)
"""A metadata index takes time to build in proportion to the table's size on disk, so on a table over its cap
the migration job must finish without it instead of running past its deadline and failing the upgrade."""
_seed_spend_logs(scratch_database, partitioned)
capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1)
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
@ -950,22 +955,20 @@ def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_b
assert f"in place except {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 ensure_request_log_indexes(scratch_database, "public", (ERROR_INDEX_DEFINITION,)) is True
assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True
@requires_db
@pytest.mark.parametrize("partitioned", (False, True), ids=("plain", "partitioned"))
def test_the_statement_logged_for_an_index_over_its_cap_builds_it_and_the_next_build_finds_it_in_place(
def test_the_statements_logged_for_an_index_over_its_cap_build_it_and_the_next_build_finds_it_in_place(
scratch_database: str, partitioned: bool, caplog: pytest.LogCaptureFixture
) -> None:
_seed_analyzed_spend_logs(scratch_database, partitioned)
capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=1)
_seed_spend_logs(scratch_database, partitioned)
capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1)
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
logged: Final = next(message for message in caplog.messages if "Not building index" in message)
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute(logged.split("finds it in place: ", 1)[1].encode())
_run_the_logged_statements(scratch_database, caplog.messages)
hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX]
caplog.clear()
@ -973,15 +976,18 @@ def test_the_statement_logged_for_an_index_over_its_cap_builds_it_and_the_next_b
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
assert "Not building index" not in caplog.text
assert _index_oids(scratch_database)[ERROR_INDEX] == hand_built
assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True
if partitioned:
_assert_index_covers_every_partition(scratch_database, ERROR_INDEX, "error_code_startTime_idx")
else:
assert _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: True}
def _leave_an_invalid_error_index(database_url: str) -> None:
with _pin_a_snapshot_on(database_url, "LiteLLM_SpendLogs"):
def _leave_an_invalid_error_index(database_url: str, name: str, table: str) -> None:
with _pin_a_snapshot_on(database_url, table):
with psycopg.connect(database_url, autocommit=True) as builder:
builder.execute("SET statement_timeout = '1s'")
prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY {} ON {} ").format(
sql.Identifier(ERROR_INDEX), sql.Identifier("LiteLLM_SpendLogs")
sql.Identifier(name), sql.Identifier(table)
)
with pytest.raises(psycopg.errors.QueryCanceled):
builder.execute(prefix.as_bytes(builder) + ERROR_INDEX_DEFINITION.body.encode())
@ -991,37 +997,63 @@ def _leave_an_invalid_error_index(database_url: str) -> None:
def test_an_interrupted_build_of_an_index_over_its_cap_is_left_to_the_logged_reindex_instead_of_repaired_at_boot(
scratch_database: str, caplog: pytest.LogCaptureFixture
) -> None:
_seed_analyzed_spend_logs(scratch_database, partitioned=False)
_leave_an_invalid_error_index(scratch_database)
capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=1)
_seed_spend_logs(scratch_database, partitioned=False)
_leave_an_invalid_error_index(scratch_database, ERROR_INDEX, "LiteLLM_SpendLogs")
capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1)
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: False}
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
logged: Final = next(message for message in caplog.messages if "Not building index" in message)
with psycopg.connect(scratch_database, autocommit=True) as conn:
conn.execute(logged.split("finds it in place: ", 1)[1].encode())
_run_the_logged_statements(scratch_database, caplog.messages)
assert _index_validity(scratch_database, "error_code_startTime_idx") == {ERROR_INDEX: True}
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
@requires_db
def test_a_table_never_analyzed_is_counted_before_deciding_whether_it_is_over_the_cap(
def test_a_partitioned_build_interrupted_between_partitions_over_its_cap_is_finished_by_the_logged_statements(
scratch_database: str, caplog: pytest.LogCaptureFixture
) -> None:
"""A partitioned build that stopped with one child attached, one left invalid and one never started must be
finished by the logged statements, and the boot-time repair must leave the invalid child to them."""
_seed_spend_logs(scratch_database, partitioned=True)
suffix: Final = "error_code_startTime_idx"
attached, interrupted = PARTITIONS
with psycopg.connect(scratch_database, autocommit=True) as conn:
for n in range(SEEDED_ROWS):
_insert_spend_log(conn, f"unanalyzed-{n}", "2026-10-01")
capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=SEEDED_ROWS - 1)
parent: Final = sql.SQL("CREATE INDEX {} ON ONLY {} ").format(
sql.Identifier(ERROR_INDEX), sql.Identifier("LiteLLM_SpendLogs")
)
conn.execute(parent.as_bytes(conn) + ERROR_INDEX_DEFINITION.body.encode())
child: Final = sql.SQL("CREATE INDEX CONCURRENTLY {} ON {} ").format(
sql.Identifier(f"{attached}_{suffix}"), sql.Identifier(attached)
)
conn.execute(child.as_bytes(conn) + ERROR_INDEX_DEFINITION.body.encode())
conn.execute(
sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format(
sql.Identifier(ERROR_INDEX), sql.Identifier(f"{attached}_{suffix}")
)
)
_leave_an_invalid_error_index(scratch_database, f"{interrupted}_{suffix}", interrupted)
capped: Final = replace(ERROR_INDEX_DEFINITION, max_bytes=1)
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
assert _index_validity(scratch_database, suffix) == {
ERROR_INDEX: False,
f"{attached}_{suffix}": True,
f"{interrupted}_{suffix}": False,
}
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
_run_the_logged_statements(scratch_database, caplog.messages)
caplog.clear()
assert f"Not building index {ERROR_INDEX}" in caplog.text
assert _indexed_table(scratch_database, ERROR_INDEX) is None
_assert_index_covers_every_partition(scratch_database, ERROR_INDEX, suffix)
with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
assert ensure_request_log_indexes(scratch_database, "public", (capped,)) is True
assert "Not building index" not in caplog.text
@requires_db

View file

@ -7,6 +7,7 @@ import re
import sqlite3
from collections.abc import Mapping
from datetime import timezone
from types import MappingProxyType
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -4772,24 +4773,25 @@ async def test_ui_view_spend_logs_with_error_code(client):
@pytest.mark.asyncio
async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client):
prefix = "E" * 128
mock_spend_logs = [
async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client: TestClient) -> None:
prefix: Final = "E" * 128
codes: Final = MappingProxyType({f"log-{suffix}": prefix + suffix for suffix in ("first", "second")})
mock_spend_logs: Final = [
{
"id": f"log-{suffix}",
"request_id": f"req-{suffix}",
"id": log_id,
"request_id": f"req-{log_id}",
"api_key": "sk-test-key",
"spend": 0.05,
"startTime": datetime.datetime.now(timezone.utc).isoformat(),
"model": "gpt-4",
"metadata": json.dumps({"error_information": {"error_code": prefix + suffix}}),
"metadata": json.dumps({"error_information": {"error_code": code}}),
}
for suffix in ("first", "second")
for log_id, code in codes.items()
]
def filter_by_exact_error_code(where):
wanted = str(where.get("metadata", {}).get("equals", "")).strip('"')
return [log for log in mock_spend_logs if json.loads(log["metadata"])["error_information"]["error_code"] == wanted]
def filter_by_exact_error_code(where: Mapping[str, Mapping[str, object]]) -> list[dict[str, str | float]]:
wanted: Final = str(where.get("metadata", {}).get("equals", "")).strip('"')
return [log for log in mock_spend_logs if codes[str(log["id"])] == wanted]
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"