mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
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:
parent
29691d6ddd
commit
2616c71da8
4 changed files with 165 additions and 104 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue