mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(spend-logs): index an error code prefix, cap the auto build at 1M rows, and leave interrupted capped builds to the operator
The error code index keys on left(code, 128) so a long error code can no longer fail a spend log write, and the filter compares that prefix next to the exact code. The index job builds it only up to 1M rows, analyzing a never-analyzed table first, logs a runnable CREATE INDEX CONCURRENTLY or REINDEX statement above that, and names the index it left in the summary. The invalid-index repair at boot skips row-capped indexes.
This commit is contained in:
parent
3bc7c5f785
commit
f4d31a5401
4 changed files with 157 additions and 33 deletions
|
|
@ -17,7 +17,7 @@ import re
|
|||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.migration_data_checks import (
|
||||
|
|
@ -75,9 +75,9 @@ REQUEST_LOG_INDEXES: Final = (
|
|||
RequestLogIndex(
|
||||
"LiteLLM_SpendLogs",
|
||||
"LiteLLM_SpendLogs_error_code_startTime_idx",
|
||||
"((metadata->'error_information'->>'error_code'), \"startTime\")",
|
||||
"((left(metadata->'error_information'->>'error_code', 128)), \"startTime\")",
|
||||
"(metadata->'error_information'->>'error_code') IS NOT NULL",
|
||||
max_rows=5_000_000,
|
||||
max_rows=1_000_000,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -94,7 +94,8 @@ _CREATE_INDEX_STATEMENT: Final = re.compile(
|
|||
)
|
||||
_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' "
|
||||
"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' "
|
||||
"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 = (
|
||||
|
|
@ -155,6 +156,10 @@ class _TableKind:
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class _EstimatedRows:
|
||||
rows: int
|
||||
unanalyzed: bool
|
||||
|
||||
|
||||
_IndexOutcome = Literal["in_place", "left_for_operator", "not_in_place"]
|
||||
|
||||
|
||||
def filter_request_log_index_diff(diff_sql: str, indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES) -> str:
|
||||
|
|
@ -198,13 +203,21 @@ def ensure_request_log_indexes(
|
|||
try:
|
||||
with connect(database_url) as connection:
|
||||
connection.execute("SET statement_timeout = 0")
|
||||
results: Final = tuple(_ensure_index(connection, schema, index) for index in indexes)
|
||||
outcomes: Final = tuple(_ensure_index(connection, schema, index) for index in indexes)
|
||||
except psycopg.Error as exc:
|
||||
logger.warning("Could not build the request-log indexes, leaving them for the next index build: %s", exc)
|
||||
return False
|
||||
if not all(results):
|
||||
if "not_in_place" in outcomes:
|
||||
logger.warning("Some request-log indexes are not in place yet, leaving them for the next index build")
|
||||
return False
|
||||
left_for_operator: Final = tuple(
|
||||
index.name for index, outcome in zip(indexes, outcomes) if outcome == "left_for_operator"
|
||||
)
|
||||
if left_for_operator:
|
||||
logger.warning(
|
||||
"Request-log indexes are in place except %s, left for an operator to build", ", ".join(left_for_operator)
|
||||
)
|
||||
return True
|
||||
logger.info("Request-log indexes are all in place")
|
||||
return True
|
||||
|
||||
|
|
@ -245,50 +258,98 @@ def _with_bounded_lock(
|
|||
return False
|
||||
|
||||
|
||||
def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool:
|
||||
def _ensure_index(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex
|
||||
) -> _IndexOutcome:
|
||||
from psycopg.rows import class_row
|
||||
|
||||
with connection.cursor(row_factory=class_row(_TableKind)) as cursor:
|
||||
table: Final = cursor.execute(_TABLE_KIND_SQL, (_regclass_name(connection, schema, index.table),)).fetchone()
|
||||
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)
|
||||
return "in_place"
|
||||
if _too_large_to_build(connection, schema, index, partitioned=table.partitioned):
|
||||
return "left_for_operator"
|
||||
built: Final = (
|
||||
build_index_on_partitioned_table(connection, schema, index)
|
||||
if table.partitioned
|
||||
else _build_leaf_index(connection, schema, index.table, index.name, index)
|
||||
)
|
||||
return "in_place" if built else "not_in_place"
|
||||
|
||||
|
||||
def _too_large_to_build(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex
|
||||
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex, partitioned: bool
|
||||
) -> 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)})
|
||||
)
|
||||
rows: Final = _estimated_rows(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",
|
||||
"and the next index build finds it in place: %s",
|
||||
index.name,
|
||||
index.table,
|
||||
rows,
|
||||
index.max_rows,
|
||||
index.body,
|
||||
_operator_build_statement(
|
||||
connection,
|
||||
schema,
|
||||
index,
|
||||
partitioned,
|
||||
interrupted=existing is not None and existing.table == index.table,
|
||||
),
|
||||
)
|
||||
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 _planner_estimate(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> _EstimatedRows:
|
||||
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)
|
||||
|
||||
|
||||
def _operator_build_statement(
|
||||
connection: "psycopg.Connection[tuple[object, ...]]",
|
||||
schema: str,
|
||||
index: RequestLogIndex,
|
||||
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."""
|
||||
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 {} "
|
||||
)
|
||||
prefix: Final = sql.SQL(build).format(sql.Identifier(index.name), sql.Identifier(schema, index.table))
|
||||
return _create_index_statement(connection, prefix, index.body).decode()
|
||||
|
||||
|
||||
def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str:
|
||||
from psycopg import sql
|
||||
|
||||
|
|
|
|||
|
|
@ -74,6 +74,9 @@ _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
|
||||
)
|
||||
_INVALID_LITELLM_INDEXES_SQL: Final = (
|
||||
"SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
|
||||
"FROM pg_index i "
|
||||
|
|
@ -851,7 +854,7 @@ class ProxyExtrasDBManager:
|
|||
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)
|
||||
return tuple(_InvalidIndex(*row) for row in rows if row[1] not in _ROW_CAPPED_REQUEST_LOG_INDEXES)
|
||||
|
||||
@staticmethod
|
||||
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
|
||||
|
|
|
|||
|
|
@ -2933,7 +2933,10 @@ async def ui_view_spend_logs(
|
|||
sql_params.append(f"%{key_alias}%")
|
||||
p += 1
|
||||
if error_code is not None:
|
||||
sql_conditions.append(f"metadata->'error_information'->>'error_code' = ${p}")
|
||||
sql_conditions.append(
|
||||
f"left(metadata->'error_information'->>'error_code', 128) = left(${p}, 128) "
|
||||
f"AND metadata->'error_information'->>'error_code' = ${p}"
|
||||
)
|
||||
sql_params.append(error_code)
|
||||
p += 1
|
||||
if error_message is not None:
|
||||
|
|
|
|||
|
|
@ -575,7 +575,8 @@ def test_partitioning_a_table_the_index_job_already_covered_moves_its_indexes_as
|
|||
_indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_litellm_call_id_idx") == "LiteLLM_SpendLogs_legacy"
|
||||
)
|
||||
assert (
|
||||
_indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_error_code_startTime_idx") == "LiteLLM_SpendLogs_legacy"
|
||||
_indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_error_code_startTime_idx")
|
||||
== "LiteLLM_SpendLogs_legacy"
|
||||
)
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
_assert_index_covers_every_partition(scratch_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
|
|
@ -946,6 +947,7 @@ def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_b
|
|||
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 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
|
||||
|
|
@ -953,18 +955,73 @@ def test_an_index_over_its_row_cap_is_left_for_an_operator_without_failing_the_b
|
|||
|
||||
|
||||
@requires_db
|
||||
def test_an_index_an_operator_built_by_hand_on_a_table_over_the_cap_is_found_in_place(
|
||||
@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(
|
||||
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)
|
||||
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())
|
||||
hand_built: Final = _index_oids(scratch_database)[ERROR_INDEX]
|
||||
caplog.clear()
|
||||
|
||||
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
|
||||
assert _index_oids(scratch_database)[ERROR_INDEX] == hand_built
|
||||
assert _index_validity(scratch_database, "error_code_startTime_idx")[ERROR_INDEX] is True
|
||||
|
||||
|
||||
def _leave_an_invalid_error_index(database_url: str) -> None:
|
||||
with _pin_a_snapshot_on(database_url, "LiteLLM_SpendLogs"):
|
||||
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")
|
||||
)
|
||||
with pytest.raises(psycopg.errors.QueryCanceled):
|
||||
builder.execute(prefix.as_bytes(builder) + ERROR_INDEX_DEFINITION.body.encode())
|
||||
|
||||
|
||||
@requires_db
|
||||
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)
|
||||
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]
|
||||
_leave_an_invalid_error_index(scratch_database)
|
||||
capped: Final = replace(ERROR_INDEX_DEFINITION, max_rows=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", (replace(ERROR_INDEX_DEFINITION, max_rows=1),))
|
||||
assert "Not building index" not in caplog.text
|
||||
assert _index_oids(scratch_database)[ERROR_INDEX] == hand_built
|
||||
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())
|
||||
|
||||
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(
|
||||
scratch_database: str, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
|
||||
@requires_db
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue