mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-09 08:18:12 +00:00
perf(spend-logs): index error rows so the Logs page error filters stop scanning every row
This commit is contained in:
mateo-berri
parent
93056ef939
commit
d0b3ecf4ad
7 changed files with 236 additions and 43 deletions
|
|
@ -57,6 +57,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
|
|||
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_litellm_call_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_error_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_error_startTime_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
|
|
|
|||
|
|
@ -5,8 +5,10 @@ background, once it serves).
|
|||
A migration cannot build them: a plain `CREATE INDEX` blocks spend-log inserts for the
|
||||
whole build, and `CREATE INDEX CONCURRENTLY` is refused on a partitioned parent
|
||||
(db_scripts/partition_spend_logs.sql). `REQUEST_LOG_INDEXES` is the one list to extend;
|
||||
names match what Prisma derives from the `@@index` declarations in schema.prisma, so an
|
||||
index a database already has is recognized and never rebuilt.
|
||||
names of plain indexes match what Prisma derives from the `@@index` declarations in
|
||||
schema.prisma, so an index a database already has is recognized and never rebuilt.
|
||||
Prisma cannot declare a partial index, and its schema diff ignores one, so a partial
|
||||
index lives only here.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
|
|
@ -33,20 +35,27 @@ if TYPE_CHECKING:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RequestLogIndex:
|
||||
"""One index the migration job owns: the table, the exact Prisma index name and the
|
||||
column list as it would be written after `ON <table>`."""
|
||||
"""One index the migration job owns: the table, the exact Prisma index name, the
|
||||
column list as it would be written after `ON <table>`, and the WHERE predicate of a
|
||||
partial index."""
|
||||
|
||||
table: str
|
||||
name: str
|
||||
definition: str
|
||||
predicate: str | None = None
|
||||
|
||||
@property
|
||||
def columns(self) -> tuple[str, ...]:
|
||||
return tuple(re.findall(r'"([^"]+)"', self.definition))
|
||||
|
||||
@property
|
||||
def body(self) -> str:
|
||||
return self.definition if self.predicate is None else f"{self.definition} WHERE {self.predicate}"
|
||||
|
||||
def partition_index_name(self, partition: str) -> str:
|
||||
"""The child index name for one partition, built the way Postgres names the
|
||||
children of a partitioned index, and kept within the 63 byte identifier limit."""
|
||||
"""The child index name this job builds on one partition, kept within the 63 byte
|
||||
identifier limit. Postgres names the children it creates for a new partition after
|
||||
their columns instead, so a partition's child is found by attachment, never by name."""
|
||||
name: Final = f"{partition}_{self.name.removeprefix(f'{self.table}_')}"
|
||||
if len(name.encode()) <= _IDENTIFIER_MAX_BYTES:
|
||||
return name
|
||||
|
|
@ -59,6 +68,12 @@ class RequestLogIndex:
|
|||
REQUEST_LOG_INDEXES: Final = (
|
||||
RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'),
|
||||
RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'),
|
||||
RequestLogIndex(
|
||||
"LiteLLM_SpendLogs",
|
||||
"LiteLLM_SpendLogs_error_startTime_idx",
|
||||
'("startTime")',
|
||||
"jsonb_typeof(metadata->'error_information') = 'object'",
|
||||
),
|
||||
)
|
||||
|
||||
_IDENTIFIER_MAX_BYTES: Final = 63
|
||||
|
|
@ -256,9 +271,12 @@ def _equivalent_indexes(
|
|||
) -> tuple[_EquivalentIndex, ...]:
|
||||
"""The indexes on `table` other than `name` with the same definition: default btree
|
||||
over the same columns in the same order, no expression, predicate, DESC or custom
|
||||
opclass or collation, and not attached under a partitioned index. Valid ones first."""
|
||||
opclass or collation, and not attached under a partitioned index. Valid ones first.
|
||||
Empty for a partial index, whose predicate this lookup cannot compare."""
|
||||
from psycopg.rows import class_row
|
||||
|
||||
if index.predicate is not None:
|
||||
return ()
|
||||
with connection.cursor(row_factory=class_row(_EquivalentIndex)) as cursor:
|
||||
return tuple(
|
||||
cursor.execute(
|
||||
|
|
@ -371,7 +389,7 @@ def _build_leaf_index(
|
|||
prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
|
||||
sql.Identifier(name), sql.Identifier(schema, table)
|
||||
)
|
||||
connection.execute(_create_index_statement(connection, prefix, index.definition))
|
||||
connection.execute(_create_index_statement(connection, prefix, index.body))
|
||||
built: Final = _index_state(connection, schema, name)
|
||||
return built is not None and built.valid
|
||||
|
||||
|
|
@ -440,7 +458,7 @@ def _create_parent_index(
|
|||
prefix: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
|
||||
sql.Identifier(name), sql.Identifier(schema, table)
|
||||
)
|
||||
statement: Final = _create_index_statement(connection, prefix, index.definition)
|
||||
statement: Final = _create_index_statement(connection, prefix, index.body)
|
||||
connection.execute(statement)
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ _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}",
|
||||
|
|
@ -2932,6 +2933,8 @@ 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)
|
||||
|
|
|
|||
|
|
@ -44,6 +44,8 @@ API_KEY_INDEX_MIGRATION: Final = "20260823000000_add_spend_logs_api_key_starttim
|
|||
CALL_ID_INDEX_MIGRATION: Final = "20260831120001_spend_logs_litellm_call_id_index"
|
||||
API_KEY_INDEX: Final = "LiteLLM_SpendLogs_api_key_startTime_idx"
|
||||
CALL_ID_INDEX: Final = "LiteLLM_SpendLogs_litellm_call_id_idx"
|
||||
ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_startTime_idx"
|
||||
START_TIME_INDEX: Final = "LiteLLM_SpendLogs_startTime_idx"
|
||||
PARTITIONED_PARENT_ERROR: Final = 'cannot create index on partitioned table "LiteLLM_SpendLogs" concurrently'
|
||||
ORIGINAL_MIGRATION_SQL: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -258,6 +260,21 @@ def _index_oids(database_url: str) -> Mapping[str, int]:
|
|||
return MappingProxyType({row.name: row.oid for row in rows})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PredicateRow:
|
||||
predicate: "str | None"
|
||||
|
||||
|
||||
def _index_predicate(database_url: str, index: str) -> "str | None":
|
||||
with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_PredicateRow)) as cursor:
|
||||
row: Final = cursor.execute(
|
||||
"SELECT pg_get_expr(indpred, indrelid) AS predicate FROM pg_index WHERE indexrelid = to_regclass(%s)",
|
||||
(f'"{index}"',),
|
||||
).fetchone()
|
||||
assert row is not None, index
|
||||
return row.predicate
|
||||
|
||||
|
||||
def _migration_job(use_v2_resolver: bool) -> bool:
|
||||
return ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver)
|
||||
|
||||
|
|
@ -287,13 +304,14 @@ def _assert_index_covers_every_partition(database_url: str, parent_index: str, s
|
|||
@requires_db
|
||||
@RESOLVERS
|
||||
@RELEASES
|
||||
def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_a_rerun_is_idempotent(
|
||||
def test_a_partitioned_spend_logs_upgrade_builds_every_index_per_partition_and_a_rerun_is_idempotent(
|
||||
partitioned_database: str, use_v2_resolver: bool
|
||||
) -> None:
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
|
||||
_assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
|
||||
_assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
|
||||
_assert_index_covers_every_partition(partitioned_database, ERROR_INDEX, "error_startTime_idx")
|
||||
ledger: Final = _assert_every_ledger_row_is_finished(partitioned_database)
|
||||
_assert_no_pending_migrations(partitioned_database)
|
||||
oids: Final = _index_oids(partitioned_database)
|
||||
|
|
@ -303,14 +321,11 @@ def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_
|
|||
'CREATE TABLE "LiteLLM_SpendLogs_p2026_10" PARTITION OF "LiteLLM_SpendLogs" '
|
||||
"FOR VALUES FROM ('2026-10-01') TO ('2026-11-01')"
|
||||
)
|
||||
inherited: Final = frozenset(
|
||||
("LiteLLM_SpendLogs_p2026_10", f"LiteLLM_SpendLogs_p2026_10_{suffix}")
|
||||
for suffix in ("api_key_startTime_idx", "litellm_call_id_idx")
|
||||
)
|
||||
attached: Final = _attached_children(partitioned_database, API_KEY_INDEX) | _attached_children(
|
||||
partitioned_database, CALL_ID_INDEX
|
||||
)
|
||||
assert inherited <= attached, attached
|
||||
inherited: Final = {
|
||||
index: "LiteLLM_SpendLogs_p2026_10" in {table for table, _ in _attached_children(partitioned_database, index)}
|
||||
for index in (API_KEY_INDEX, CALL_ID_INDEX, ERROR_INDEX)
|
||||
}
|
||||
assert all(inherited.values()), inherited
|
||||
|
||||
assert _migration_job(use_v2_resolver) is True
|
||||
assert _ledger(partitioned_database) == ledger
|
||||
|
|
@ -320,7 +335,7 @@ def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_
|
|||
@requires_db
|
||||
@RESOLVERS
|
||||
@RELEASES
|
||||
def test_a_plain_spend_logs_upgrade_builds_both_indexes_and_a_second_job_run_rebuilds_nothing(
|
||||
def test_a_plain_spend_logs_upgrade_builds_every_index_and_a_second_job_run_rebuilds_nothing(
|
||||
scratch_database: str, use_v2_resolver: bool
|
||||
) -> None:
|
||||
with psycopg.connect(scratch_database, autocommit=True) as conn:
|
||||
|
|
@ -331,6 +346,7 @@ def test_a_plain_spend_logs_upgrade_builds_both_indexes_and_a_second_job_run_reb
|
|||
|
||||
assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
|
||||
assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
|
||||
assert _index_validity(scratch_database, "error_startTime_idx") == {ERROR_INDEX: True}
|
||||
_assert_every_ledger_row_is_finished(scratch_database)
|
||||
_assert_no_pending_migrations(scratch_database)
|
||||
oids: Final = _index_oids(scratch_database)
|
||||
|
|
@ -358,7 +374,8 @@ def test_a_database_that_applied_the_original_migration_files_sees_no_pending_mi
|
|||
assert _migration_job(use_v2_resolver) is True
|
||||
|
||||
assert _ledger(scratch_database) == before
|
||||
assert _index_oids(scratch_database) == oids
|
||||
after: Final = _index_oids(scratch_database)
|
||||
assert {name: after[name] for name in oids} == oids
|
||||
|
||||
|
||||
@requires_db
|
||||
|
|
@ -539,22 +556,28 @@ def test_an_invalid_child_index_left_by_an_interrupted_build_is_rebuilt_and_atta
|
|||
|
||||
|
||||
@requires_db
|
||||
def test_partitioning_a_table_the_index_job_already_covered_moves_its_call_id_index_aside(
|
||||
def test_partitioning_a_table_the_index_job_already_covered_moves_its_indexes_aside(
|
||||
scratch_database: str,
|
||||
) -> None:
|
||||
"""The runbook renames the Prisma indexes off the legacy table so the partitioned table can take their
|
||||
names. The call-id index is built by the index job, not Prisma, and a deployment that partitions after
|
||||
that job ran kept the name on the legacy table, so every later job run reported it as on another table."""
|
||||
names. The call-id and error indexes are built by the index job, not Prisma, and a deployment that partitions
|
||||
after that job ran kept their names on the legacy table, so every later job run reported them as on another
|
||||
table."""
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
assert _indexed_table(scratch_database, CALL_ID_INDEX) == "LiteLLM_SpendLogs"
|
||||
assert _indexed_table(scratch_database, ERROR_INDEX) == "LiteLLM_SpendLogs"
|
||||
|
||||
_partition_spend_logs(scratch_database)
|
||||
|
||||
assert (
|
||||
_indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_litellm_call_id_idx") == "LiteLLM_SpendLogs_legacy"
|
||||
)
|
||||
assert (
|
||||
_indexed_table(scratch_database, "LiteLLM_SpendLogs_legacy_error_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")
|
||||
_assert_index_covers_every_partition(scratch_database, ERROR_INDEX, "error_startTime_idx")
|
||||
|
||||
|
||||
@requires_db
|
||||
|
|
@ -882,6 +905,21 @@ def test_a_hand_built_child_index_under_another_name_is_renamed_and_attached(par
|
|||
assert "p2026_08_call_id_by_hand" not in oids and oids[f"{partition}_litellm_call_id_idx"] == hand_built
|
||||
|
||||
|
||||
@requires_db
|
||||
def test_the_partial_error_index_is_built_beside_the_plain_start_time_index_never_renamed_from_it(
|
||||
scratch_database: str,
|
||||
) -> None:
|
||||
start_time_index: Final = _index_oids(scratch_database)[START_TIME_INDEX]
|
||||
|
||||
assert ensure_request_log_indexes(scratch_database, "public") is True
|
||||
|
||||
oids: Final = _index_oids(scratch_database)
|
||||
assert oids[START_TIME_INDEX] == start_time_index
|
||||
assert oids[ERROR_INDEX] != start_time_index
|
||||
assert _index_predicate(scratch_database, START_TIME_INDEX) is None
|
||||
assert _index_predicate(scratch_database, ERROR_INDEX) is not None
|
||||
|
||||
|
||||
@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:
|
||||
|
|
@ -921,7 +959,7 @@ def test_a_valid_partitioned_parent_index_under_another_name_is_renamed_with_its
|
|||
with caplog.at_level("INFO", logger="litellm_proxy_extras"):
|
||||
assert ensure_request_log_indexes(partitioned_database, "public") is True
|
||||
|
||||
assert "Building index" not in caplog.text
|
||||
assert not [line for line in caplog.messages if line.startswith("Building index") and "litellm_call_id" in line]
|
||||
oids: Final = _index_oids(partitioned_database)
|
||||
assert "call_id_parent_by_hand" not in oids and oids[CALL_ID_INDEX] == hand_built
|
||||
assert _attached_children(partitioned_database, CALL_ID_INDEX) == children_before
|
||||
|
|
|
|||
|
|
@ -43,7 +43,8 @@ MIGRATION_JOB_SECONDS: Final = 300
|
|||
INDEXES_IN_PLACE: Final = "Request-log indexes are all in place"
|
||||
INDEX_BUILD_LINES: Final = ("Building index", "Attached index")
|
||||
BUILD_SECONDS: Final = 60
|
||||
SPEND_LOGS_INDEXES: Final = ("LiteLLM_SpendLogs_api_key_startTime_idx", "LiteLLM_SpendLogs_litellm_call_id_idx")
|
||||
ORIGINAL_SPEND_LOGS_INDEXES: Final = ("LiteLLM_SpendLogs_api_key_startTime_idx", "LiteLLM_SpendLogs_litellm_call_id_idx")
|
||||
SPEND_LOGS_INDEXES: Final = (*ORIGINAL_SPEND_LOGS_INDEXES, "LiteLLM_SpendLogs_error_startTime_idx")
|
||||
POPULATED_PARTITIONS: Final = MappingProxyType(
|
||||
{
|
||||
"LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"),
|
||||
|
|
@ -254,9 +255,9 @@ def assert_ready(booted_gateway: Gateway) -> None:
|
|||
assert readiness.json()["db"] == "connected", readiness.text
|
||||
|
||||
|
||||
def assert_both_indexes_cover_every_partition(database_url: str) -> None:
|
||||
"""Every populated partition's index is attached and valid, both parents are valid, and
|
||||
a partition created afterwards inherits both indexes."""
|
||||
def assert_every_index_covers_every_partition(database_url: str) -> None:
|
||||
"""Every populated partition's index is attached and valid, every parent is valid, and
|
||||
a partition created afterwards inherits every index."""
|
||||
populated: Final = (*POPULATED_PARTITIONS, DEFAULT_PARTITION)
|
||||
assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
|
||||
assert attached_partition_indexes(database_url) == expected_attachments(populated)
|
||||
|
|
@ -285,7 +286,7 @@ def assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(
|
|||
|
||||
|
||||
@RESOLVERS
|
||||
def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_both_indexes_per_partition(
|
||||
def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_every_index_per_partition(
|
||||
gateway: Gateway, tmp_path: Path, resolver: Resolver
|
||||
) -> None:
|
||||
with scratch_database() as database_url:
|
||||
|
|
@ -294,7 +295,7 @@ def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_bot
|
|||
job: Final = migration_job(database_url, resolver)
|
||||
assert job.returncode == 0, job.stdout + job.stderr
|
||||
assert INDEXES_IN_PLACE in job.stderr + job.stdout, job.stdout + job.stderr
|
||||
assert_both_indexes_cover_every_partition(database_url)
|
||||
assert_every_index_covers_every_partition(database_url)
|
||||
assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
|
||||
|
||||
|
||||
|
|
@ -308,7 +309,7 @@ def test_the_migration_job_heals_a_partitioned_table_left_with_the_failed_call_i
|
|||
fail_the_call_id_index_migration_like_the_shipped_release(database_url, tmp_path)
|
||||
job: Final = migration_job(database_url, resolver)
|
||||
assert job.returncode == 0, job.stdout + job.stderr
|
||||
assert_both_indexes_cover_every_partition(database_url)
|
||||
assert_every_index_covers_every_partition(database_url)
|
||||
assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
|
||||
|
||||
|
||||
|
|
@ -321,23 +322,25 @@ def test_the_migration_job_leaves_a_plain_table_that_applied_the_original_index_
|
|||
deployed: Final = deploy_the_original_index_migrations(database_url, tmp_path)
|
||||
assert deployed.returncode == 0, deployed.stdout + deployed.stderr
|
||||
before: Final = index_oids(database_url)
|
||||
assert set(SPEND_LOGS_INDEXES) <= set(before), before
|
||||
assert set(ORIGINAL_SPEND_LOGS_INDEXES) <= set(before), before
|
||||
job: Final = migration_job(database_url, resolver)
|
||||
assert job.returncode == 0, job.stdout + job.stderr
|
||||
assert INDEXES_IN_PLACE in job.stderr + job.stdout, job.stdout + job.stderr
|
||||
assert "Building index" not in job.stderr + job.stdout, job.stdout + job.stderr
|
||||
output: Final = job.stdout + job.stderr
|
||||
assert job.returncode == 0, output
|
||||
assert INDEXES_IN_PLACE in output, output
|
||||
assert not any(f"Building index {name} " in output for name in ORIGINAL_SPEND_LOGS_INDEXES), output
|
||||
assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
|
||||
assert index_oids(database_url) == before
|
||||
after: Final = index_oids(database_url)
|
||||
assert {name: after[name] for name in before} == before
|
||||
assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
|
||||
|
||||
|
||||
@RESOLVERS
|
||||
def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_indexes_after_it_is_ready(
|
||||
def test_a_serving_proxy_that_runs_the_migrations_itself_builds_every_index_after_it_is_ready(
|
||||
gateway: Gateway, tmp_path: Path, resolver: Resolver
|
||||
) -> None:
|
||||
"""A deployment that runs migrate deploy from the serving proxy and never runs the
|
||||
migration job answers readiness with the inert files applied, then its background build
|
||||
puts both indexes on every partition."""
|
||||
puts every index on every partition."""
|
||||
with scratch_database() as database_url:
|
||||
deploy_schema_before_the_index_migrations(database_url, tmp_path)
|
||||
partition_spend_logs(database_url)
|
||||
|
|
@ -353,13 +356,13 @@ def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_indexes_aft
|
|||
seconds=BUILD_SECONDS,
|
||||
)
|
||||
assert "Building index" in log and "Attached index" in log, log[-4000:]
|
||||
assert_both_indexes_cover_every_partition(database_url)
|
||||
assert_every_index_covers_every_partition(database_url)
|
||||
|
||||
|
||||
def test_the_cli_run_as_a_migration_job_builds_both_indexes_before_it_exits(gateway: Gateway, tmp_path: Path) -> None:
|
||||
def test_the_cli_run_as_a_migration_job_builds_every_index_before_it_exits(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with scratch_database() as database_url:
|
||||
deploy_schema_before_the_index_migrations(database_url, tmp_path)
|
||||
partition_spend_logs(database_url)
|
||||
job: Final = migration_cli(database_url, gateway, V2)
|
||||
assert job.returncode == 0, job.stdout + job.stderr
|
||||
assert_both_indexes_cover_every_partition(database_url)
|
||||
assert_every_index_covers_every_partition(database_url)
|
||||
|
|
|
|||
129
tests/integration/spend/test_spend_log_error_filters.py
Normal file
129
tests/integration/spend/test_spend_log_error_filters.py
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, list_value, object_value, string_value
|
||||
from integration._support.database import scratch_database
|
||||
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)
|
||||
|
||||
ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_startTime_idx"
|
||||
INDEXES_IN_PLACE: Final = "Request-log indexes are all in place"
|
||||
SUCCESS_ROWS: Final = 5000
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ErrorRow:
|
||||
request_id: str
|
||||
status: str
|
||||
error_code: str
|
||||
error_message: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ScanRow:
|
||||
scans: int
|
||||
|
||||
|
||||
ERROR_ROWS: Final = (
|
||||
ErrorRow("failed-500-overloaded", "failure", "500", "upstream overloaded"),
|
||||
ErrorRow("failed-500-timeout", "failure", "500", "upstream timed out"),
|
||||
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"})),
|
||||
)
|
||||
SESSION_VIEWS: Final = ("false", "true")
|
||||
|
||||
|
||||
def _seed(database_url: str, started: datetime) -> None:
|
||||
utc: Final = started.astimezone(timezone.utc).replace(tzinfo=None)
|
||||
with psycopg.connect(database_url, autocommit=True) as connection:
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime", metadata, status) '
|
||||
"SELECT 'ok-' || n, 'acompletion', %(started)s + n * interval '1 millisecond', "
|
||||
"%(started)s + n * interval '1 millisecond', %(metadata)s, 'success' "
|
||||
"FROM generate_series(1, %(rows)s) AS n",
|
||||
{"started": utc, "metadata": Jsonb({"error_information": None}), "rows": SUCCESS_ROWS},
|
||||
)
|
||||
with connection.cursor() as cursor:
|
||||
cursor.executemany(
|
||||
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime", metadata, status) '
|
||||
"VALUES (%s, 'acompletion', %s, %s, %s, %s)",
|
||||
tuple(
|
||||
(
|
||||
row.request_id,
|
||||
utc,
|
||||
utc,
|
||||
Jsonb(
|
||||
{
|
||||
"error_information": {
|
||||
"error_code": row.error_code,
|
||||
"error_message": row.error_message,
|
||||
"error_class": "APIError",
|
||||
}
|
||||
}
|
||||
),
|
||||
row.status,
|
||||
)
|
||||
for row in ERROR_ROWS
|
||||
),
|
||||
)
|
||||
connection.execute('ANALYZE "LiteLLM_SpendLogs"')
|
||||
|
||||
|
||||
def _error_index_scans(database_url: str) -> int:
|
||||
with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_ScanRow)) as cursor:
|
||||
row: Final = cursor.execute(
|
||||
"SELECT idx_scan AS scans FROM pg_stat_user_indexes WHERE indexrelname = %s", (ERROR_INDEX,)
|
||||
).fetchone()
|
||||
assert row is not None, f"{ERROR_INDEX} was not built"
|
||||
return row.scans
|
||||
|
||||
|
||||
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:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
window: Final = {
|
||||
"start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
}
|
||||
with scratch_database() as database_url:
|
||||
with owned_proxy_process(
|
||||
gateway, tmp_path, {"DATABASE_URL": database_url}, database_setup=MIGRATE_DEPLOY
|
||||
) as booted:
|
||||
eventually(
|
||||
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,
|
||||
)
|
||||
|
|
@ -34,10 +34,10 @@ def _prisma_indexes_of(schema: str, model: str) -> frozenset[str]:
|
|||
|
||||
|
||||
class TestTheIndexList:
|
||||
def test_every_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self):
|
||||
def test_every_plain_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self):
|
||||
schema: Final = SCHEMA.read_text()
|
||||
for index in REQUEST_LOG_INDEXES:
|
||||
assert index.name in _prisma_indexes_of(schema, index.table), index
|
||||
assert (index.predicate is None) == (index.name in _prisma_indexes_of(schema, index.table)), index
|
||||
|
||||
@pytest.mark.parametrize("name", INERT_MIGRATIONS)
|
||||
def test_the_migrations_that_used_to_build_these_indexes_run_no_sql(self, name: str):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue