mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
264 lines
11 KiB
Python
264 lines
11 KiB
Python
import os
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Iterator, Mapping
|
|
from types import MappingProxyType
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager
|
|
|
|
psycopg = pytest.importorskip("psycopg")
|
|
|
|
pytestmark = pytest.mark.timeout(120)
|
|
|
|
requires_db: Final = pytest.mark.skipif(
|
|
"DATABASE_URL" not in os.environ,
|
|
reason="requires a postgres database (DATABASE_URL)",
|
|
)
|
|
|
|
HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable"
|
|
HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx"
|
|
HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC'
|
|
LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable"
|
|
LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx"
|
|
PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable"
|
|
PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx"
|
|
|
|
|
|
def _base_url() -> str:
|
|
return os.environ["DATABASE_URL"].split("?")[0]
|
|
|
|
|
|
def _index_validity(schema: str) -> Mapping[str, bool]:
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
rows = conn.execute(
|
|
"SELECT c.relname, i.indisvalid FROM pg_index i "
|
|
"JOIN pg_class c ON c.oid = i.indexrelid "
|
|
"JOIN pg_namespace n ON n.oid = c.relnamespace "
|
|
"WHERE n.nspname = %s",
|
|
(schema,),
|
|
).fetchall()
|
|
return MappingProxyType(dict(rows))
|
|
|
|
|
|
def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None:
|
|
"""Abort a CONCURRENTLY build while it waits on an older snapshot, the same
|
|
spot the deadlock loser dies at, so it leaves its index INVALID."""
|
|
with psycopg.connect(_base_url()) as pin:
|
|
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
|
pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
|
|
with psycopg.connect(_base_url(), autocommit=True) as builder:
|
|
builder.execute("SET statement_timeout = '1s'")
|
|
with pytest.raises(psycopg.errors.QueryCanceled):
|
|
builder.execute(statement)
|
|
|
|
|
|
def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None:
|
|
_interrupt_concurrent_build(
|
|
schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})'
|
|
)
|
|
|
|
|
|
def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None:
|
|
_interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"')
|
|
|
|
|
|
@pytest.fixture
|
|
def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
|
|
schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'CREATE SCHEMA "{schema}"')
|
|
conn.execute(
|
|
f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)'
|
|
)
|
|
conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)')
|
|
|
|
monkeypatch.delenv("DIRECT_URL", raising=False)
|
|
monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}")
|
|
yield schema
|
|
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'DROP SCHEMA "{schema}" CASCADE')
|
|
|
|
|
|
@pytest.fixture
|
|
def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
|
|
"""A brand-new database, what a first deploy sees. A scratch schema would
|
|
not do: the migrations guard on pg_constraint by name across every schema,
|
|
so a LiteLLM schema already pushed into public makes them skip and then
|
|
fail, which is exactly what CI's database looks like."""
|
|
admin_url: Final = _base_url()
|
|
name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}"
|
|
with psycopg.connect(admin_url, autocommit=True) as conn:
|
|
conn.execute(f'CREATE DATABASE "{name}"')
|
|
|
|
monkeypatch.delenv("DIRECT_URL", raising=False)
|
|
monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}")
|
|
yield "public"
|
|
|
|
with psycopg.connect(admin_url, autocommit=True) as conn:
|
|
conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)')
|
|
|
|
|
|
@requires_db
|
|
def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None:
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
_leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id")
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False}
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False}
|
|
|
|
|
|
@requires_db
|
|
def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None:
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
_leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX)
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"')
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"')
|
|
before: Final = _index_validity(scratch_schema)
|
|
assert len(before) == 4
|
|
assert set(before.values()) == {False}
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
|
|
|
|
|
@requires_db
|
|
def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None:
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})')
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
|
|
|
|
|
@requires_db
|
|
def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None:
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)')
|
|
conn.execute(
|
|
f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" '
|
|
f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)'
|
|
)
|
|
conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)')
|
|
assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
|
|
assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False}
|
|
|
|
|
|
@requires_db
|
|
def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None:
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
|
|
with psycopg.connect(_base_url(), autocommit=True) as other_replica:
|
|
other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,))
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: False}
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
|
|
|
|
|
@requires_db
|
|
def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None:
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
|
|
with psycopg.connect(_base_url()) as pin:
|
|
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
|
pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"')
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False
|
|
blocked: Final = _index_validity(scratch_schema)
|
|
assert blocked[HEALTH_INDEX] is False
|
|
assert [name for name in blocked if name.endswith("_ccnew")]
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
|
|
|
|
|
def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None:
|
|
with psycopg.connect(_base_url()) as pin:
|
|
pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
|
|
pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"')
|
|
pinned.set()
|
|
pin.execute("SELECT pg_sleep(%s)", (seconds,))
|
|
|
|
|
|
@requires_db
|
|
def test_repair_outlives_a_statement_timeout_passed_through_database_url_options(
|
|
scratch_schema: str, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000")
|
|
pinned: Final = threading.Event()
|
|
holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0))
|
|
holder.start()
|
|
pinned.wait()
|
|
try:
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
finally:
|
|
holder.join()
|
|
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
|
|
|
|
|
@requires_db
|
|
def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}"
|
|
index: Final = f"{table}_id_idx"
|
|
monkeypatch.setenv("DATABASE_URL", _base_url())
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'CREATE TABLE public."{table}" (id TEXT)')
|
|
try:
|
|
_leave_invalid_index("public", table, index, "id")
|
|
assert _index_validity("public")[index] is False
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
|
|
assert _index_validity("public")[index] is True
|
|
finally:
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'DROP TABLE public."{table}"')
|
|
|
|
|
|
def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("DIRECT_URL", raising=False)
|
|
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever")
|
|
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is False
|
|
|
|
|
|
@requires_db
|
|
def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None:
|
|
_leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
with pytest.MonkeyPatch.context() as env:
|
|
env.setenv("DIRECT_URL", f"{_base_url()}?schema=public")
|
|
env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}")
|
|
assert ProxyExtrasDBManager.repair_invalid_indexes() is True
|
|
|
|
assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
|
|
|
|
|
|
def _invalidate_deployed_index(schema: str) -> None:
|
|
with psycopg.connect(_base_url(), autocommit=True) as conn:
|
|
conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"')
|
|
_leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
|
|
|
|
|
|
@requires_db
|
|
@pytest.mark.timeout(300)
|
|
@pytest.mark.parametrize("use_v2_resolver", [True, False])
|
|
def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None:
|
|
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
|
|
_invalidate_deployed_index(fresh_database)
|
|
assert _index_validity(fresh_database)[HEALTH_INDEX] is False
|
|
|
|
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
|
|
|
|
assert _index_validity(fresh_database)[HEALTH_INDEX] is True
|