From 55c7872496b5568bc259fc88c0645f702af1576e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:57:06 -0700 Subject: [PATCH 01/16] fix(proxy-extras): rebuild indexes left INVALID by a migration deadlock Two replicas racing prisma migrate deploy can deadlock, and the loser dies mid CREATE INDEX CONCURRENTLY, leaving the index INVALID. The retried migration's IF NOT EXISTS then skips it, so the planner never uses it. After migrations succeed, look for INVALID indexes on LiteLLM tables and have one replica (advisory try-lock) REINDEX INDEX CONCURRENTLY each of them, dropping _ccnew/_ccold leftovers of an interrupted rebuild instead. The repair never blocks startup: a failed rebuild is logged and retried on the next boot. Also encode DATABASE_URL query values with quote instead of quote_plus so options=-c%20... reaches psycopg intact. --- .../litellm_proxy_extras/utils.py | 128 +++++++++- .../test_invalid_index_repair.py | 238 ++++++++++++++++++ 2 files changed, 359 insertions(+), 7 deletions(-) create mode 100644 tests/proxy_migration_tests/test_invalid_index_repair.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..df019d78078 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -6,18 +6,23 @@ import shutil import subprocess import tempfile import time +from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import TYPE_CHECKING, Final, Optional from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.replica_identity import ( - REPLICA_IDENTITY_FULL_ENV_VAR, - apply_replica_identity_full, -) from litellm_proxy_extras.prisma_toolchain import ( ensure_prisma_toolchain, prisma_command_timeout, ) +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) + +if TYPE_CHECKING: + import psycopg + import psycopg.sql def str_to_bool(value: Optional[str]) -> bool: @@ -40,6 +45,29 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") +INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") +_INVALID_LITELLM_INDEXES_SQL: Final = ( + "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " + "FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_class t ON t.oid = i.indrelid " + "JOIN pg_namespace n ON n.oid = t.relnamespace " + "WHERE NOT i.indisvalid " + " AND c.relkind = 'i' " + " AND n.nspname = %s " + " AND t.relname LIKE %s " + " AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) " + "ORDER BY c.relname" +) + + +@dataclass(frozen=True, slots=True) +class _InvalidIndex: + schema: str + name: str + table_size: str + _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile( r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE @@ -557,7 +585,7 @@ class ProxyExtrasDBManager: def _strip_prisma_query_params(url: str) -> str: """Remove Prisma-specific query params (connection_limit, pool_timeout, schema, etc.) from DATABASE_URL so psycopg can parse it.""" - from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode + from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse parsed = urlparse(url) if not parsed.query: @@ -578,7 +606,7 @@ class ProxyExtrasDBManager: "target_session_attrs", } kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params] - return urlunparse(parsed._replace(query=urlencode(kept))) + return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote))) @staticmethod def _warn_if_db_ahead_of_head(migrations_dir: str) -> None: @@ -652,6 +680,91 @@ class ProxyExtrasDBManager: ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), ) + @staticmethod + 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) + + @staticmethod + def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]: + from psycopg import sql + + target: Final = sql.Identifier(index.schema, index.name) + if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name): + return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover" + return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt" + + @staticmethod + def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None: + import psycopg + + statement, action = ProxyExtrasDBManager._index_repair(index) + try: + conn.execute(statement) + except psycopg.Error as e: + logger.warning( + "Could not repair invalid index %s.%s, will retry on the next startup. " + "If this keeps happening, run `%s` by hand as the index owner. Error: %s", + index.schema, + index.name, + statement.as_string(conn), + e, + ) + return + logger.info("%s invalid index %s.%s", action, index.schema, index.name) + + @staticmethod + def repair_invalid_indexes(lock_timeout: str = "30s") -> bool: + """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left + INVALID (a migration deadlock between replicas is the usual cause; the + retried migration skips them because of IF NOT EXISTS). Never raises: + returns True when no invalid index remains, False when the repair was + skipped or failed and will be retried on the next startup.""" + database_url: Final = os.getenv("DATABASE_URL") + if not database_url: + return False + + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the invalid index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + conn.execute("SET statement_timeout = 0") + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + if not found: + return True + logger.warning( + "Found %d invalid index(es) left by an interrupted CREATE INDEX " + "CONCURRENTLY, rebuilding: %s", + len(found), + ", ".join(f"{index.name} (table size {index.table_size})" for index in found), + ) + lock_row: Final = conn.execute( + "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,) + ).fetchone() + if lock_row is None or not lock_row[0]: + logger.info("Another replica is already rebuilding the invalid indexes, skipping") + return False + for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema): + ProxyExtrasDBManager._repair_index(conn, index) + remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + except psycopg.Error as e: + logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e) + return False + return not remaining + @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ @@ -886,6 +999,7 @@ class ProxyExtrasDBManager: use_migrate=use_migrate, use_v2_resolver=use_v2_resolver ) if migrated: + ProxyExtrasDBManager.repair_invalid_indexes() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return migrated diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py new file mode 100644 index 00000000000..787a62b361b --- /dev/null +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -0,0 +1,238 @@ +import os +import subprocess +import threading +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +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.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') + + +@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.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + + +class _MigrateDeployApplied: + stdout = "Applied migration.\n" + stderr = "" + returncode = 0 + + +@requires_db +@pytest.mark.parametrize("use_v2_resolver", [True, False]) +def test_setup_database_repairs_the_index_after_a_recovered_deploy( + scratch_schema: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, use_v2_resolver: bool +) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", lambda *_, **__: True) + monkeypatch.setattr(subprocess, "run", lambda *_, **__: _MigrateDeployApplied()) + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} From 2d64020d432f1c6112b836a752720a25a5ab93e9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:24:42 -0700 Subject: [PATCH 02/16] fix(proxy-extras): run the index repair over DIRECT_URL and exercise it through a real migrate deploy --- .../litellm_proxy_extras/utils.py | 7 ++- .../test_invalid_index_repair.py | 63 ++++++++++++------- 2 files changed, 46 insertions(+), 24 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index df019d78078..75fda59f28e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -721,8 +721,11 @@ class ProxyExtrasDBManager: INVALID (a migration deadlock between replicas is the usual cause; the retried migration skips them because of IF NOT EXISTS). Never raises: returns True when no invalid index remains, False when the repair was - skipped or failed and will be retried on the next startup.""" - database_url: Final = os.getenv("DATABASE_URL") + skipped or failed and will be retried on the next startup. Runs over + DIRECT_URL when set: the session settings, the advisory lock and REINDEX + CONCURRENTLY all need one server session, which a transaction pooler + does not give.""" + database_url: Final = os.getenv("DIRECT_URL") or os.getenv("DATABASE_URL") if not database_url: return False diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py index 787a62b361b..2e46775cd2b 100644 --- a/tests/proxy_migration_tests/test_invalid_index_repair.py +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -1,9 +1,7 @@ import os -import subprocess import threading import uuid from collections.abc import Iterator, Mapping -from pathlib import Path from types import MappingProxyType from typing import Final @@ -66,16 +64,14 @@ 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]: +def _scratch_schema(monkeypatch: pytest.MonkeyPatch, *table_definitions: str) -> 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)') + for definition in table_definitions: + conn.execute(f'CREATE TABLE "{schema}".{definition}') + monkeypatch.delenv("DIRECT_URL", raising=False) monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") yield schema @@ -83,6 +79,20 @@ def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: conn.execute(f'DROP SCHEMA "{schema}" CASCADE') +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + yield from _scratch_schema( + monkeypatch, + f'"{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)', + f'"{LOOKALIKE_TABLE}" (id TEXT)', + ) + + +@pytest.fixture +def empty_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + yield from _scratch_schema(monkeypatch) + + @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) @@ -210,29 +220,38 @@ def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) - 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 -class _MigrateDeployApplied: - stdout = "Applied migration.\n" - stderr = "" - returncode = 0 +@requires_db +def test_repair_runs_over_direct_url_when_set(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + direct_url: Final = os.environ["DATABASE_URL"] + with pytest.MonkeyPatch.context() as env: + env.setenv("DIRECT_URL", direct_url) + env.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + 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( - scratch_schema: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, use_v2_resolver: bool -) -> None: - _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) - (tmp_path / "schema.prisma").write_text("// stub") - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", lambda *_, **__: True) - monkeypatch.setattr(subprocess, "run", lambda *_, **__: _MigrateDeployApplied()) +def test_setup_database_repairs_the_index_after_a_recovered_deploy(empty_schema: str, use_v2_resolver: bool) -> None: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + _invalidate_deployed_index(empty_schema) + assert _index_validity(empty_schema)[HEALTH_INDEX] is False assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True - assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + assert _index_validity(empty_schema)[HEALTH_INDEX] is True From bfdaedf51b13f4bc8872d51a99d7f288fc0cc56f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:23:50 -0700 Subject: [PATCH 03/16] fix(proxy-extras): take the repair schema from DATABASE_URL and deploy the real-migration test on a fresh database --- .../litellm_proxy_extras/utils.py | 17 +++---- .../test_invalid_index_repair.py | 47 +++++++++++-------- 2 files changed, 36 insertions(+), 28 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 75fda59f28e..40bc4cd1dfd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -721,12 +721,13 @@ class ProxyExtrasDBManager: INVALID (a migration deadlock between replicas is the usual cause; the retried migration skips them because of IF NOT EXISTS). Never raises: returns True when no invalid index remains, False when the repair was - skipped or failed and will be retried on the next startup. Runs over - DIRECT_URL when set: the session settings, the advisory lock and REINDEX - CONCURRENTLY all need one server session, which a transaction pooler - does not give.""" - database_url: Final = os.getenv("DIRECT_URL") or os.getenv("DATABASE_URL") - if not database_url: + skipped or failed and will be retried on the next startup. Looks in the + schema DATABASE_URL names, the only URL Prisma migrates through, but + connects over DIRECT_URL when set: the session settings, the advisory + lock and REINDEX CONCURRENTLY all need one server session, which a + transaction pooler does not give.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: return False try: @@ -739,8 +740,8 @@ class ProxyExtrasDBManager: ) return False - schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" - cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) try: with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: conn.execute("SET statement_timeout = 0") diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py index 2e46775cd2b..741fa7386df 100644 --- a/tests/proxy_migration_tests/test_invalid_index_repair.py +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -64,12 +64,15 @@ def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') -def _scratch_schema(monkeypatch: pytest.MonkeyPatch, *table_definitions: str) -> Iterator[str]: +@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}"') - for definition in table_definitions: - conn.execute(f'CREATE TABLE "{schema}".{definition}') + 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}") @@ -80,17 +83,22 @@ def _scratch_schema(monkeypatch: pytest.MonkeyPatch, *table_definitions: str) -> @pytest.fixture -def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: - yield from _scratch_schema( - monkeypatch, - f'"{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)', - f'"{LOOKALIKE_TABLE}" (id TEXT)', - ) +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" -@pytest.fixture -def empty_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: - yield from _scratch_schema(monkeypatch) + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)') @requires_db @@ -227,12 +235,11 @@ def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch @requires_db -def test_repair_runs_over_direct_url_when_set(scratch_schema: str) -> None: +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) - direct_url: Final = os.environ["DATABASE_URL"] with pytest.MonkeyPatch.context() as env: - env.setenv("DIRECT_URL", direct_url) - env.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + 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} @@ -247,11 +254,11 @@ def _invalidate_deployed_index(schema: str) -> None: @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(empty_schema: str, use_v2_resolver: bool) -> None: +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(empty_schema) - assert _index_validity(empty_schema)[HEALTH_INDEX] is False + _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(empty_schema)[HEALTH_INDEX] is True + assert _index_validity(fresh_database)[HEALTH_INDEX] is True From 48cc4efca38b46e1f561da4728954d06bbe223c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:37:17 -0700 Subject: [PATCH 04/16] fix(least-busy): share in-flight request counts across proxy workers Least-busy kept one dict of in-flight counts per model group in the router cache, which reads in-memory first, so every worker and replica routed on its own stale copy and each write overwrote the shared value. Counts now live in one key per deployment, incremented and read through Redis when the router has a Redis cache, and in the process-local cache otherwise. --- basedpyright-code-budget.json | 12 +- litellm/caching/redis_cache.py | 4 +- litellm/router_strategy/least_busy.py | 372 +++++++++--------- ruff-strict-budget.json | 10 +- .../local_testing/test_least_busy_routing.py | 46 ++- .../router_strategy/test_least_busy.py | 151 +++++++ type-discipline-budget.json | 8 +- 7 files changed, 376 insertions(+), 227 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_least_busy.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..c21dc76743a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5570 + "limit": 5551 }, "reportMissingTypeArgument": { - "limit": 15281 + "limit": 15277 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44358 + "limit": 44357 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38240 }, "reportUnknownParameterType": { - "limit": 19584 + "limit": 19558 }, "reportUnknownVariableType": { - "limit": 29814 + "limit": 29781 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index aaee7188d86..d4c47914394 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -680,7 +680,7 @@ class RedisCache(BaseCache): # NON blocking - notify users Redis is throwing an exception print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") - def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: + def increment_cache(self, key, value: int, ttl: float | None = None, refresh_ttl: bool = False, **kwargs) -> int: _redis_client: Final = self.redis_client start_time = time.time() set_ttl: Final = self.get_ttl(ttl=ttl) @@ -701,7 +701,7 @@ class RedisCache(BaseCache): if set_ttl is not None: # check if key already has ttl, if not -> set ttl start_time = time.time() - current_ttl: Final = _redis_client.ttl(key) + current_ttl: Final = -1 if refresh_ttl else _redis_client.ttl(key) end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..de6a4c4f59a 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,17 +1,96 @@ -#### What this does #### -# identifies least busy deployment -# How is this achieved? -# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"} -# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic} -# - use litellm.success + failure callbacks to log when a request completed -# - in get_available_deployment, for a given model group name -> pick based on traffic - -import random +from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 + + +class _ModelInfo(TypedDict, total=False): + id: ReadOnly[str | int | None] + + +class _Metadata(TypedDict, total=False): + model_group: ReadOnly[str | None] + + +class _LitellmParams(TypedDict, total=False): + metadata: ReadOnly[_Metadata | None] + model_info: ReadOnly[_ModelInfo | None] + + +class _CallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[_LitellmParams | None] + + +class _DeploymentModelInfo(TypedDict): + id: ReadOnly[str | int] + + +class _Deployment(TypedDict): + model_info: ReadOnly[_DeploymentModelInfo] + + +_CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) +_REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...]) + + +def _request_count_key(model_group: str, deployment_id: str) -> str: + return f"{model_group}_request_count:{deployment_id}" + + +def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: + try: + call: Final = _CALL_KWARGS.validate_python(kwargs) + except ValidationError: + return None + litellm_params: Final = call.get("litellm_params") + metadata: Final = litellm_params.get("metadata") if litellm_params else None + model_info: Final = litellm_params.get("model_info") if litellm_params else None + model_group: Final = metadata.get("model_group") if metadata else None + deployment_id: Final = model_info.get("id") if model_info else None + if model_group is None or deployment_id is None: + return None + return model_group, str(deployment_id) + + +def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + return tuple( + _request_count_key(model_group, str(deployment["model_info"]["id"])) + for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + ) + + +def _as_count(value: float | None) -> int: + return 0 if value is None else int(value) + + +def _least_busy( + healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] +) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] + + +def _warn_unreadable(model_group: str, error: Exception) -> None: + verbose_router_logger.warning( + "least-busy routing could not read the in-flight counts for %s, treating every deployment as idle: %s", + model_group, + error, + ) + + +def _warn_unwritable(key: str, error: Exception) -> None: + verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error) + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False @@ -21,194 +100,101 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache - def log_pre_api_call(self, model, messages, kwargs): - """ - Log when a model is being used. + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + self._increment(kwargs, 1) - Caching based on model group. - """ - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - request_count_api_key: Final = f"{model_group}_request_count" - # update cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - except Exception: - pass + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - def log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - def _get_available_deployments( - self, - healthy_deployments: list, - all_deployments: dict, - ): - """ - Helper to get deployments using least busy strategy - """ - for d in healthy_deployments: - ## if healthy deployment not yet used - if d["model_info"]["id"] not in all_deployments: - all_deployments[d["model_info"]["id"]] = 0 - # map deployment to id - # pick least busy deployment - min_traffic = float("inf") - min_deployment = None - for k, v in all_deployments.items(): - if v < min_traffic: - min_traffic = v - min_deployment = k - if min_deployment is not None: - ## check if min deployment is a string, if so, cast it to int - for m in healthy_deployments: - if m["model_info"]["id"] == min_deployment: - return m - min_deployment = random.choice(healthy_deployments) - else: - min_deployment = random.choice(healthy_deployments) - return min_deployment + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 def get_available_deployments( - self, - model_group: str, - healthy_deployments: list, - ): - """ - Sync helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + try: + counts: Final = tuple(_as_count(value) for value in self._read_counts(keys)) + except Exception as e: + _warn_unreadable(model_group, e) + return _least_busy(healthy_deployments, (0,) * len(keys)) + return _least_busy(healthy_deployments, counts) - async def async_get_available_deployments(self, model_group: str, healthy_deployments: list): - """ - Async helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + async def async_get_available_deployments( + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + try: + counts: Final = tuple(_as_count(value) for value in await self._async_read_counts(keys)) + except Exception as e: + _warn_unreadable(model_group, e) + return _least_busy(healthy_deployments, (0,) * len(keys)) + return _least_busy(healthy_deployments, counts) + + def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + if redis_cache is None: + self.router_cache.increment_cache(key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + else: + redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + except Exception as e: + _warn_unwritable(key, e) + + async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + if redis_cache is None: + await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + else: + await redis_cache.async_increment(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + except Exception as e: + _warn_unwritable(key, e) + + def _read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return _MEMORY_COUNTS.validate_python(self.router_cache.batch_get_cache(list(keys), local_only=True)) + by_key: Final = _REDIS_COUNTS.validate_python(redis_cache.batch_get_cache(key_list=list(keys))) + return tuple(by_key.get(key) for key in keys) + + async def _async_read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return _MEMORY_COUNTS.validate_python( + await self.router_cache.async_batch_get_cache(list(keys), local_only=True) + ) + by_key: Final = _REDIS_COUNTS.validate_python(await redis_cache.async_batch_get_cache(key_list=list(keys))) + return tuple(by_key.get(key) for key in keys) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..ceff8c6e15c 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2956 + "limit": 2937 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1979 + "limit": 1972 }, "ANN202": { - "limit": 831 + "limit": 830 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2915 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 217 + "limit": 212 }, "S112": { "limit": 22 diff --git a/tests/local_testing/test_least_busy_routing.py b/tests/local_testing/test_least_busy_routing.py index 18ab8bf779d..fb83d4e601f 100644 --- a/tests/local_testing/test_least_busy_routing.py +++ b/tests/local_testing/test_least_busy_routing.py @@ -33,8 +33,8 @@ def test_model_added(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"gpt-3.5-turbo_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = "gpt-3.5-turbo_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 def test_get_available_deployments(): @@ -52,8 +52,8 @@ def test_get_available_deployments(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"{model_group}_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = f"{model_group}_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 # test_get_available_deployments() @@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test): router.leastbusy_logger.test_flag = True model_group = "azure-model" - request_count_dict = {1: 10, 2: 54, 3: 100} - cache_key = f"{model_group}_request_count" + request_count_dict = {"1": 10, "2": 54, "3": 100} + cache_keys = { + deployment_id: f"{model_group}_request_count:{deployment_id}" + for deployment_id in request_count_dict + } if async_test is True: - await router.cache.async_set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count) deployment = await router.async_get_available_deployment( model=model_group, messages=None, request_kwargs={} ) else: - router.cache.set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + router.cache.set_cache(key=cache_keys[deployment_id], value=count) deployment = router.get_available_deployment(model=model_group, messages=None) print(f"deployment: {deployment}") assert deployment["model_info"]["id"] == "1" @@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - return_dict = router.cache.get_cache(key=cache_key) - # wait 2 seconds time.sleep(2) + return_dict = { + deployment_id: router.cache.get_cache(key=cache_key) + for deployment_id, cache_key in cache_keys.items() + } + assert router.leastbusy_logger.logged_success == 1 - assert return_dict[1] == 10 - assert return_dict[2] == 54 - assert return_dict[3] == 100 + assert return_dict["1"] == 10 + assert return_dict["2"] == 54 + assert return_dict["3"] == 100 ## Test with Real calls ## @@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.atext_completion(model=model, prompt=prompt, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" @@ -259,8 +269,10 @@ async def test_router_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.acompletion(model=model, messages=messages, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py new file mode 100644 index 00000000000..32a226a5d5b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -0,0 +1,151 @@ +import json +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler + +GROUP: Final = "least-busy-group" +DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} +DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} +HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] + + +def _call_kwargs(deployment_id: str) -> dict[str, object]: + return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} + + +class SharedRedisCounters: + """Stores JSON strings and hands back a fresh object per read, the way a real Redis client does.""" + + def __init__(self) -> None: + self.encoded: dict[str, str] = {} + self.ttls: dict[str, float] = {} + + def count(self, key: str) -> object: + raw: Final = self.encoded.get(key) + return None if raw is None else json.loads(raw) + + def get_cache(self, key: str, **kwargs: object) -> object: + return self.count(key) + + def set_cache(self, key: str, value: object, **kwargs: object) -> None: + self.encoded[key] = json.dumps(value) + + async def async_get_cache(self, key: str, **kwargs: object) -> object: + return self.count(key) + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + self.set_cache(key, value) + + def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + current: Final = self.count(key) or 0 + assert isinstance(current, int) + incremented: Final = current + value + self.encoded[key] = json.dumps(incremented) + if ttl is not None and (refresh_ttl or key not in self.ttls): + self.ttls[key] = ttl + return incremented + + async def async_increment(self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False) -> float: + return self.increment_cache(key, int(value), ttl, refresh_ttl) + + def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + return {key: self.count(key) for key in key_list} + + async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + return self.batch_get_cache(key_list) + + +def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double + return LeastBusyLoggingHandler(router_cache=cache) + + +@pytest.mark.asyncio +async def test_worker_routes_around_a_request_another_worker_started() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_sync_pick_reads_the_shared_counts() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_redis_counts_keep_a_refreshed_ttl() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + assert shared.ttls == {f"{GROUP}_request_count:dep-a": IN_FLIGHT_COUNT_TTL_SECONDS} + + +@pytest.mark.asyncio +async def test_counts_stay_in_memory_without_redis() -> None: + worker: Final = _worker(None) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + +class UnavailableRedis(SharedRedisCounters): + def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + raise ConnectionError("redis is down") + + def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + raise ConnectionError("redis is down") + + +def test_redis_outage_never_fails_the_request() -> None: + worker: Final = _worker(UnavailableRedis()) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_calls_without_a_deployment_are_ignored() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) + worker.log_pre_api_call(model="m", messages=[], kwargs={}) + + assert shared.encoded == {} diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..e413e6db2c6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22180 + "limit": 22176 }, "LIT002": { - "limit": 26729 + "limit": 26721 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16412 }, "LIT011": { - "limit": 5506 + "limit": 5505 }, "LIT012": { "limit": 4486 From c5aa4f07185c91f7c055e614734d5e216c6d043d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:14:58 -0700 Subject: [PATCH 05/16] fix(least-busy): fall back to per-worker counts when Redis is unreadable Keep each worker's own in-flight counter up to date alongside the shared one, so a Redis outage routes on that worker's counts the way it did before this branch instead of treating every deployment as idle. Floor a counter at zero when a decrement finds the key gone, which happens when a request outlives the 1 hour TTL, so an expired counter cannot settle at -1 and win every pick. --- litellm/router_strategy/least_busy.py | 97 +++++++++++-------- .../router_strategy/test_least_busy.py | 39 +++++++- 2 files changed, 95 insertions(+), 41 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index de6a4c4f59a..00d27b5f8d9 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -39,7 +39,7 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) _REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) -_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: @@ -68,8 +68,20 @@ def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[ ) -def _as_count(value: float | None) -> int: - return 0 if value is None else int(value) +def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: + return tuple(0 if value is None else int(value) for value in values) + + +def _shared_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + by_key: Final = _REDIS_COUNTS.validate_python(raw) + return _as_counts([by_key.get(key) for key in keys]) + + +def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + values: Final = _MEMORY_COUNTS.validate_python(raw) + if values is None or len(values) != len(keys): + return (0,) * len(keys) + return _as_counts(values) def _least_busy( @@ -82,7 +94,8 @@ def _least_busy( def _warn_unreadable(model_group: str, error: Exception) -> None: verbose_router_logger.warning( - "least-busy routing could not read the in-flight counts for %s, treating every deployment as idle: %s", + "least-busy routing could not read the shared in-flight counts for %s, " + "falling back to this worker's own counts: %s", model_group, error, ) @@ -135,23 +148,31 @@ class LeastBusyLoggingHandler(CustomLogger): self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> Mapping[str, object] | None: keys: Final = _request_count_keys(model_group, healthy_deployments) - try: - counts: Final = tuple(_as_count(value) for value in self._read_counts(keys)) - except Exception as e: - _warn_unreadable(model_group, e) - return _least_busy(healthy_deployments, (0,) * len(keys)) - return _least_busy(healthy_deployments, counts) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _shared_counts(redis_cache.batch_get_cache(key_list=list(keys)), keys) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) async def async_get_available_deployments( self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> Mapping[str, object] | None: keys: Final = _request_count_keys(model_group, healthy_deployments) - try: - counts: Final = tuple(_as_count(value) for value in await self._async_read_counts(keys)) - except Exception as e: - _warn_unreadable(model_group, e) - return _least_busy(healthy_deployments, (0,) * len(keys)) - return _least_busy(healthy_deployments, counts) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _shared_counts(await redis_cache.async_batch_get_cache(key_list=list(keys)), keys) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: ref: Final = _deployment_ref(kwargs) @@ -160,10 +181,16 @@ class LeastBusyLoggingHandler(CustomLogger): key: Final = _request_count_key(*ref) redis_cache: Final = self.router_cache.redis_cache try: + local: Final = self.router_cache.increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local < 0: + self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: - self.router_cache.increment_cache(key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - else: - redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + return + shared: Final = redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + if shared < 0: + redis_cache.set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) @@ -174,27 +201,17 @@ class LeastBusyLoggingHandler(CustomLogger): key: Final = _request_count_key(*ref) redis_cache: Final = self.router_cache.redis_cache try: + local: Final = await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local is not None and local < 0: + await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: - await self.router_cache.async_increment_cache( - key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS - ) - else: - await redis_cache.async_increment(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + return + shared: Final = await redis_cache.async_increment( + key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True + ) + if shared < 0: + await redis_cache.async_set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) - - def _read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: - redis_cache: Final = self.router_cache.redis_cache - if redis_cache is None: - return _MEMORY_COUNTS.validate_python(self.router_cache.batch_get_cache(list(keys), local_only=True)) - by_key: Final = _REDIS_COUNTS.validate_python(redis_cache.batch_get_cache(key_list=list(keys))) - return tuple(by_key.get(key) for key in keys) - - async def _async_read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: - redis_cache: Final = self.router_cache.redis_cache - if redis_cache is None: - return _MEMORY_COUNTS.validate_python( - await self.router_cache.async_batch_get_cache(list(keys), local_only=True) - ) - by_key: Final = _REDIS_COUNTS.validate_python(await redis_cache.async_batch_get_cache(key_list=list(keys))) - return tuple(by_key.get(key) for key in keys) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index 32a226a5d5b..55702c6fe73 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -133,14 +133,51 @@ class UnavailableRedis(SharedRedisCounters): raise ConnectionError("redis is down") -def test_redis_outage_never_fails_the_request() -> None: +@pytest.mark.asyncio +async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: worker: Final = _worker(UnavailableRedis()) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + shared.encoded.clear() + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.asyncio +async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None: + worker: Final = _worker(None) + in_memory: Final = worker.router_cache.in_memory_cache + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + in_memory.delete_cache(f"{GROUP}_request_count:dep-a") + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + def test_calls_without_a_deployment_are_ignored() -> None: shared: Final = SharedRedisCounters() worker: Final = _worker(shared) From a9bc2cb50b4ba443e72b7d03a33447cc8724b037 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:38:26 -0700 Subject: [PATCH 06/16] fix(least-busy): clamp the shared in-flight count to zero in one call A decrement whose matching increment is gone, because the counter key expired while the request was still in flight, used to recreate the key at -1, and a deployment with a negative count looks permanently idle, so it collects every pick from then on. The repair write that followed the decrement could also land after another pod's increment and erase it. The increment, the clamp at zero and the TTL refresh now run as a single Lua call, so nothing can interleave between them. --- litellm/caching/redis_cache.py | 39 ++++++++++++++++++- litellm/router_strategy/least_busy.py | 10 +---- .../router_strategy/test_least_busy.py | 20 +++++----- 3 files changed, 50 insertions(+), 19 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d4c47914394..4eccde5742b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,6 +20,8 @@ from contextvars import ContextVar from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -80,11 +82,22 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} ) +_INCREMENT_WITH_FLOOR_LUA: Final = ( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " + "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " + "redis.call('EXPIRE', KEYS[1], ARGV[2]) " + "return count" +) + +_LUA_COUNT: Final = TypeAdapter(int) + def _get_call_stack_info(num_frames: int = 2) -> str: """ @@ -680,7 +693,7 @@ class RedisCache(BaseCache): # NON blocking - notify users Redis is throwing an exception print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") - def increment_cache(self, key, value: int, ttl: float | None = None, refresh_ttl: bool = False, **kwargs) -> int: + def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: _redis_client: Final = self.redis_client start_time = time.time() set_ttl: Final = self.get_ttl(ttl=ttl) @@ -701,7 +714,7 @@ class RedisCache(BaseCache): if set_ttl is not None: # check if key already has ttl, if not -> set ttl start_time = time.time() - current_ttl: Final = -1 if refresh_ttl else _redis_client.ttl(key) + current_ttl: Final = _redis_client.ttl(key) end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( @@ -736,6 +749,20 @@ class RedisCache(BaseCache): ) raise e + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Add ``value`` to ``key``, clamp the result at zero, and refresh the TTL, in one Lua call. + + A counter whose key expired while a request was still in flight would otherwise be + recreated negative by that request's decrement. Clamping inside the same call is what + keeps it safe: a separate corrective write could land after another pod's increment and + erase it. Returns the resulting count. + """ + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl + ) + return _LUA_COUNT.validate_python(count) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() @@ -1241,6 +1268,14 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) + return _LUA_COUNT.validate_python(count) + async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 00d27b5f8d9..771d2bb4328 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -188,9 +188,7 @@ class LeastBusyLoggingHandler(CustomLogger): self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: return - shared: Final = redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) - if shared < 0: - redis_cache.set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) @@ -208,10 +206,6 @@ class LeastBusyLoggingHandler(CustomLogger): await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: return - shared: Final = await redis_cache.async_increment( - key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True - ) - if shared < 0: - await redis_cache.async_set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index 55702c6fe73..eb9591811ac 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -40,17 +40,16 @@ class SharedRedisCounters: async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: self.set_cache(key, value) - def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: current: Final = self.count(key) or 0 assert isinstance(current, int) - incremented: Final = current + value + incremented: Final = max(0, current + value) self.encoded[key] = json.dumps(incremented) - if ttl is not None and (refresh_ttl or key not in self.ttls): - self.ttls[key] = ttl + self.ttls[key] = ttl return incremented - async def async_increment(self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False) -> float: - return self.increment_cache(key, int(value), ttl, refresh_ttl) + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + return self.increment_with_floor(key, value, ttl) def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: return {key: self.count(key) for key in key_list} @@ -102,12 +101,14 @@ def test_sync_pick_reads_the_shared_counts() -> None: def test_redis_counts_keep_a_refreshed_ttl() -> None: shared: Final = SharedRedisCounters() worker: Final = _worker(shared) + key: Final = f"{GROUP}_request_count:dep-a" + shared.ttls[key] = 5 worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(f"{GROUP}_request_count:dep-a") == 0 - assert shared.ttls == {f"{GROUP}_request_count:dep-a": IN_FLIGHT_COUNT_TTL_SECONDS} + assert shared.count(key) == 0 + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} @pytest.mark.asyncio @@ -129,7 +130,7 @@ class UnavailableRedis(SharedRedisCounters): def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: raise ConnectionError("redis is down") - def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: raise ConnectionError("redis is down") @@ -159,6 +160,7 @@ def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B From caf9bbbd5a535e61da844e1a0225a4f8a6312b8d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:20:05 -0700 Subject: [PATCH 07/16] fix(least-busy): keep the shared count readable, counted once, and off the loop A Redis outage read as "every deployment is idle", because batch_get_cache swallows the failure and answers with an empty dict. batch_get_counts and its async twin raise instead, so a worker that cannot reach Redis falls back to its own numbers rather than routing on zeros. The counter's TTL is now set only on a key that has none, so a +1 left behind by a worker that died mid-request ages out an hour after the key was created. It used to be refreshed on every touch, which kept that stuck count alive for as long as the group took traffic. Two least-busy groups counted the same request twice, since the pre-call list kept a selector per group while the success list deduped by class. The selector now goes on through add_litellm_input_callback, which dedupes the same way. A prompt-management model picked its deployment on the synchronous path, so the new Redis read landed on the event loop and configured routing plugins never ran. It awaits the async selector now. --- basedpyright-code-budget.json | 12 ++-- litellm/caching/redis_cache.py | 36 +++++++++- litellm/router.py | 8 ++- litellm/router_strategy/least_busy.py | 10 +-- ruff-strict-budget.json | 10 +-- .../test_litellm/caching/test_redis_cache.py | 44 ++++++++++++ .../router_strategy/test_least_busy.py | 67 +++++++++---------- .../test_router_routing_groups.py | 38 +++++++++++ .../test_router_routing_plugins.py | 30 +++++++++ type-discipline-budget.json | 8 +-- 10 files changed, 199 insertions(+), 64 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c21dc76743a..32d3730aa32 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5551 + "limit": 5532 }, "reportMissingTypeArgument": { - "limit": 15277 + "limit": 15273 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38240 + "limit": 38208 }, "reportUnknownParameterType": { - "limit": 19558 + "limit": 19532 }, "reportUnknownVariableType": { - "limit": 29781 + "limit": 29751 }, "reportUnnecessaryCast": { "limit": 110 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 138 + "limit": 137 }, "reportUnusedImport": { "limit": 542 diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 4eccde5742b..3abc6e6f3c9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -92,11 +92,18 @@ _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( _INCREMENT_WITH_FLOOR_LUA: Final = ( "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " - "redis.call('EXPIRE', KEYS[1], ARGV[2]) " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " "return count" ) _LUA_COUNT: Final = TypeAdapter(int) +_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) + + +def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: + return _OPTIONAL_COUNTS.validate_python( + tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) + ) def _get_call_stack_info(num_frames: int = 2) -> str: @@ -749,13 +756,19 @@ class RedisCache(BaseCache): ) raise e + @_redis_circuit_breaker_guard_sync def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - """Add ``value`` to ``key``, clamp the result at zero, and refresh the TTL, in one Lua call. + """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. A counter whose key expired while a request was still in flight would otherwise be recreated negative by that request's decrement. Clamping inside the same call is what keeps it safe: a separate corrective write could land after another pod's increment and - erase it. Returns the resulting count. + erase it. + + The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was + created rather than ``ttl`` after it was last touched. Refreshing it on every touch + would keep a count a dead worker never decremented alive for as long as the group + takes traffic. Returns the resulting count. """ namespaced_key: Final = self.check_and_fix_namespace(key=key) count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval @@ -763,6 +776,23 @@ class RedisCache(BaseCache): ) return _LUA_COUNT.validate_python(count) + @_redis_circuit_breaker_guard_sync + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + + ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller + cannot tell apart from "every counter is unset". A caller that has to fall back to its + own numbers when Redis is unreachable needs the failure, not a dict of zeros. + """ + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + + @_redis_circuit_breaker_guard + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Async twin of ``batch_get_counts``, raising on failure the same way.""" + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() diff --git a/litellm/router.py b/litellm/router.py index 3f450661946..f5b7924fb58 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1212,7 +1212,7 @@ class Router: selector = LeastBusyLoggingHandler(router_cache=self.cache) if register_callbacks: if isinstance(litellm.input_callback, list): - litellm.input_callback.append(selector) + litellm.logging_callback_manager.add_litellm_input_callback(selector) else: litellm.input_callback = [selector] case RoutingStrategy.USAGE_BASED_ROUTING.value: @@ -4139,10 +4139,12 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - prompt_management_deployment: Final = self.get_available_deployment( + specific_deployment: Final = kwargs.pop("specific_deployment", None) + prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), + specific_deployment=specific_deployment, + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 771d2bb4328..ab6d702bbc2 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -38,7 +38,6 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) -_REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) _MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) @@ -72,11 +71,6 @@ def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: return tuple(0 if value is None else int(value) for value in values) -def _shared_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: - by_key: Final = _REDIS_COUNTS.validate_python(raw) - return _as_counts([by_key.get(key) for key in keys]) - - def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: values: Final = _MEMORY_COUNTS.validate_python(raw) if values is None or len(values) != len(keys): @@ -151,7 +145,7 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _shared_counts(redis_cache.batch_get_cache(key_list=list(keys)), keys) + shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) except Exception as e: _warn_unreadable(model_group, e) else: @@ -166,7 +160,7 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _shared_counts(await redis_cache.async_batch_get_cache(key_list=list(keys)), keys) + shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) except Exception as e: _warn_unreadable(model_group, e) else: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ceff8c6e15c..d63a69de76f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2937 + "limit": 2918 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1972 + "limit": 1965 }, "ANN202": { - "limit": 830 + "limit": 829 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2915 + "limit": 2914 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 212 + "limit": 207 }, "S112": { "limit": 22 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2f412e7382b..6b2df118611 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} +def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): + """A caller that must fall back when Redis is unreachable needs the failure, not zeros. + + The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell + apart from "every counter is unset". Least-busy routing read that as an idle deployment and + kept sending traffic to it instead of falling back to this worker's own in-flight counts. + """ + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + sync_batch_redis_cache.batch_get_counts(["lit7039"]) + + +@pytest.mark.asyncio +async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): + """Async twin: the async batch read hides the same failure behind an empty dict.""" + failing_client = AsyncMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + with patch.object(cache, "init_async_client", return_value=failing_client): + assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_batch_get_counts(["lit7039"]) + + +@pytest.mark.parametrize("stored", [b"3", "3"]) +def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): + """Counters come back positionally, so an unset key has to stay a hole rather than shift the + rest of the row onto the wrong deployments, and a count has to survive whether the client + hands it back as bytes or as text.""" + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.return_value = [stored, None, b"0"] + + assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) + + @pytest.fixture def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: service_logger = ServiceLogging(mock_testing=True) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index eb9591811ac..9efa526fc02 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -1,4 +1,3 @@ -import json from typing import Final import pytest @@ -18,44 +17,34 @@ def _call_kwargs(deployment_id: str) -> dict[str, object]: class SharedRedisCounters: - """Stores JSON strings and hands back a fresh object per read, the way a real Redis client does.""" + """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when + the key is created, and ordered reads that raise rather than invent a value.""" def __init__(self) -> None: - self.encoded: dict[str, str] = {} - self.ttls: dict[str, float] = {} + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} - def count(self, key: str) -> object: - raw: Final = self.encoded.get(key) - return None if raw is None else json.loads(raw) + def count(self, key: str) -> int | None: + return self.counts.get(key) - def get_cache(self, key: str, **kwargs: object) -> object: - return self.count(key) - - def set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.encoded[key] = json.dumps(value) - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - return self.count(key) - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.set_cache(key, value) + def expire(self, key: str) -> None: + self.counts.pop(key, None) + self.ttls.pop(key, None) def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - current: Final = self.count(key) or 0 - assert isinstance(current, int) - incremented: Final = max(0, current + value) - self.encoded[key] = json.dumps(incremented) - self.ttls[key] = ttl + incremented: Final = max(0, self.counts.get(key, 0) + value) + self.counts[key] = incremented + self.ttls.setdefault(key, ttl) return incremented async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: return self.increment_with_floor(key, value, ttl) - def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: - return {key: self.count(key) for key in key_list} + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return tuple(self.counts.get(key) for key in key_list) - async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: - return self.batch_get_cache(key_list) + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return self.batch_get_counts(key_list) def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: @@ -98,17 +87,25 @@ def test_sync_pick_reads_the_shared_counts() -> None: assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A -def test_redis_counts_keep_a_refreshed_ttl() -> None: +def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: + """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that + stuck count an hour after the key was created, which only works while nothing writes the TTL + again: a handler that refreshed it on every touch would keep the count alive for as long as + the group takes traffic, and the deployment would read busier than it is forever.""" shared: Final = SharedRedisCounters() worker: Final = _worker(shared) key: Final = f"{GROUP}_request_count:dep-a" - shared.ttls[key] = 5 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + + shared.ttls[key] = 5 worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(key) == 0 - assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + assert shared.count(key) == 1 + assert shared.ttls == {key: 5} @pytest.mark.asyncio @@ -127,10 +124,10 @@ async def test_counts_stay_in_memory_without_redis() -> None: class UnavailableRedis(SharedRedisCounters): - def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: raise ConnectionError("redis is down") - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: raise ConnectionError("redis is down") @@ -153,7 +150,7 @@ def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: worker: Final = _worker(shared) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - shared.encoded.clear() + shared.expire(f"{GROUP}_request_count:dep-a") worker.log_success_event(_call_kwargs("dep-a"), None, None, None) assert shared.count(f"{GROUP}_request_count:dep-a") == 0 @@ -187,4 +184,4 @@ def test_calls_without_a_deployment_are_ignored() -> None: worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) worker.log_pre_api_call(model="m", messages=[], kwargs={}) - assert shared.encoded == {} + assert shared.counts == {} diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5599c5aad63..0ce0ed5b37e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy @@ -435,6 +436,43 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp assert router._group_selectors == {} +def test_two_least_busy_groups_count_a_request_once(monkeypatch): + """ + Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and + back down from the success hooks on `litellm.callbacks`. The success list drops a second + selector of the same class, so a pre-call list that kept both counted every request twice + and released it once, and the deployment's in-flight count climbed until it looked pinned. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + router = _build_router( + routing_strategy="least-busy", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "least-busy", + } + ], + ) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index 293af36080a..c49b22ea367 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -164,6 +164,36 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) +@pytest.mark.asyncio +async def test_prompt_management_model_still_runs_the_plugin_pipeline(): + """ + A prompt-management model routes through its own factory, which picked the deployment + on the synchronous path. Plugins never run there, so the guard turned every such request + into an error message about the caller's own API choice, on an async call the caller made + correctly. It also read the in-flight counts with a blocking call inside the event loop. + """ + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="lit-7039", + ) + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e413e6db2c6..0c0952289e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22176 + "limit": 22174 }, "LIT002": { - "limit": 26721 + "limit": 26715 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16412 + "limit": 16398 }, "LIT011": { - "limit": 5505 + "limit": 5504 }, "LIT012": { "limit": 4486 From ef5f51abca77d3081d04796ce9b3191dc2409f9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:59:15 -0700 Subject: [PATCH 08/16] fix(least-busy): count for every router, and keep the expiry through a clamp Two routers in one process shared a single handler, because the callback manager dedupes on the class name plus the handler's public attributes and the handler had none. The second router's requests were never counted. The handler now carries the id of the cache it was built on, so routers with different caches both register while the two selectors one router builds for its routing groups still collapse into one. Clamping a negative count back to zero used SET, which drops the key's TTL, so the next write started the hour over. It uses INCRBY by the negative amount now, which leaves the expiry alone. The Lua script had no test that ran it, so tests/local_testing covers both the sync and async paths against a real Redis, and the file is wired into the CircleCI job that provides one. --- .circleci/config.yml | 1 + litellm/caching/redis_cache.py | 2 +- litellm/router_strategy/least_busy.py | 1 + .../test_redis_increment_with_floor.py | 80 +++++++++++++++++++ .../test_router_routing_groups.py | 38 +++++++++ 5 files changed, 121 insertions(+), 1 deletion(-) create mode 100644 tests/local_testing/test_redis_increment_with_floor.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 6e368a3debe..d10962864bf 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1440,6 +1440,7 @@ jobs: TEST_FILES=$(printf "%s\n" \ tests/local_testing/test_dual_cache.py \ tests/local_testing/test_redis_batch_optimizations.py \ + tests/local_testing/test_redis_increment_with_floor.py \ tests/local_testing/test_router_utils.py) echo "$TEST_FILES" | circleci tests run \ --verbose \ diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3abc6e6f3c9..106c1580110 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -91,7 +91,7 @@ _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( _INCREMENT_WITH_FLOOR_LUA: Final = ( "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " - "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end " "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " "return count" ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index ab6d702bbc2..14e6592e1fd 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -106,6 +106,7 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache + self.router_cache_id = str(id(router_cache)) def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: self._increment(kwargs, 1) diff --git a/tests/local_testing/test_redis_increment_with_floor.py b/tests/local_testing/test_redis_increment_with_floor.py new file mode 100644 index 00000000000..e358d5f31e0 --- /dev/null +++ b/tests/local_testing/test_redis_increment_with_floor.py @@ -0,0 +1,80 @@ +"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the +create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so +these are the only tests that fail when the script itself is wrong.""" + +import os +import uuid +from typing import Final + +import pytest +from dotenv import load_dotenv + +load_dotenv() + +from litellm.caching.redis_cache import RedisCache + +TTL: Final = 600 + + +@pytest.fixture +def counter(): + cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + key: Final = f"lit7039-{uuid.uuid4()}" + yield cache, key, cache.check_and_fix_namespace(key=key) + cache.delete_cache(key) + + +def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter): + cache, key, _ = counter + + assert cache.increment_with_floor(key, 3, TTL) == 3 + assert cache.increment_with_floor(key, 2, TTL) == 5 + assert cache.batch_get_counts([key]) == (5,) + + +def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter): + """A worker whose counter expired mid-request decrements a key that is no longer there. + Without the clamp that deployment reads negative, and least-busy pins every later request + on it until the count climbs back to zero.""" + cache, key, _ = counter + + assert cache.increment_with_floor(key, 1, TTL) == 1 + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.batch_get_counts([key]) == (0,) + + +def test_traffic_never_pushes_a_counters_expiry_back_out(counter): + """The TTL is what releases a count whose worker died mid-request. Rewriting it on every + touch would keep that stuck count alive for as long as the group takes traffic.""" + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + assert cache.redis_client.ttl(namespaced_key) > TTL - 60 + + cache.redis_client.expire(namespaced_key, 30) + cache.increment_with_floor(key, 1, TTL) + + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter): + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + cache.redis_client.expire(namespaced_key, 30) + + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +@pytest.mark.asyncio +async def test_the_async_counter_behaves_the_same_way(counter): + cache, key, namespaced_key = counter + + assert await cache.async_increment_with_floor(key, 2, TTL) == 2 + assert await cache.async_batch_get_counts([key]) == (2,) + + cache.redis_client.expire(namespaced_key, 30) + + assert await cache.async_increment_with_floor(key, -9, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 0ce0ed5b37e..af390c3292b 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -473,6 +473,44 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 +def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): + """ + Least-busy hangs its counting off litellm's global callback lists, and those lists keep one + logger per class unless the instances differ in a plain attribute. Two routers in one process + (a second Router, or a per-request `user_config` one) therefore have to register separately: + a second router whose selector is dropped counts nothing, reads zero for every deployment, + and sends every request to whichever one is listed first. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + first = _build_router(routing_strategy="least-busy") + second = _build_router(routing_strategy="least-busy") + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1 + assert ( + second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"] + == "deploy-2" + ) + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- From 314e8905c583792d08207cdcd1bfa8e22d985f51 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:44:45 -0700 Subject: [PATCH 09/16] fix(router): let prompt-management plugins see the caller's own messages The prompt-management factory picks its deployment with a placeholder message. That was inert while the pick ran on the synchronous path, which never runs the routing plugin pipeline. Now that the pick runs the pipeline, a plugin classifying request content would score the placeholder instead of the conversation, and the narrowing it writes decides which deployments the real call may use. --- litellm/router.py | 2 +- .../test_router_routing_plugins.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index f5b7924fb58..398d87ce843 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4142,7 +4142,7 @@ class Router: specific_deployment: Final = kwargs.pop("specific_deployment", None) prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, - messages=[{"role": "user", "content": "prompt"}], + messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally specific_deployment=specific_deployment, request_kwargs=kwargs, ) diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index c49b22ea367..0a54addce5e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -56,6 +56,18 @@ class BlockEverything: return context +class MessageRecorder: + """Records what each plugin pass was handed, then blocks so the request stops there.""" + + def __init__(self): + self.seen = [] + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.raw_messages)) + context.candidate_models = [] + return context + + def _smart_router_model_list(): return [ { @@ -194,6 +206,41 @@ async def test_prompt_management_model_still_runs_the_plugin_pipeline(): ) +@pytest.mark.asyncio +async def test_prompt_management_plugins_see_the_callers_own_messages(): + """ + The prompt-management factory picks its deployment with a placeholder message, which was + harmless while that pick ran on the synchronous path (plugins never ran there at all). Now + that the pick runs the plugin pipeline, a plugin that classifies request content would score + the placeholder instead of the conversation, and the narrowing it produces decides which + deployments the real call is allowed to use. + """ + recorder = MessageRecorder() + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[recorder], + ) + messages = [{"role": "user", "content": "wire me $40,000 to account 12345"}] + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=messages, + litellm_call_id="lit-7039", + ) + + assert recorder.seen == [messages] + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" From 95c0f9db7d92cefafa5e2d9799786b4931c9a390 Mon Sep 17 00:00:00 2001 From: Clement Date: Wed, 9 Sep 2026 01:05:23 +0800 Subject: [PATCH 10/16] fix(router): rank streaming latency routing by raw TTFT, not TTFT per token (#40202) * fix(router): rank streaming latency routing by raw TTFT, not TTFT per token Latency-based routing divided time-to-first-token by completion_tokens before storing it, so a deployment that streamed a long answer looked faster to first token than one that answered briefly. TTFT is now stored as plain seconds (first token time minus request start) in both the sync and async success handlers, which is what the routing decision compares. Non-streaming latency normalization per output token is unchanged. Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj * fix(router): store streaming TTFT under a seconds-only cache key Workers on the previous release keep writing seconds-per-token samples under "time_to_first_token" in the shared router cache during a rolling deploy, so mixing the new raw-seconds samples into the same list averaged incompatible units. Raw TTFT now lives under "time_to_first_token_seconds" and routing reads only that key. Also fix the regression test's token counts: with 50 tokens on the fast deployment and 500 on the slow one the old per-token formula picks the slow deployment, so the routing assertion now catches the bug. Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj * test(router): cover the TTFT sliding window from the unit-test shard Move the TTFT list trimming checks from the CircleCI-only suite into the mapped unit test file as one sync/async parametrized test, so the changed lines in lowest_latency.py are exercised by the GitHub unit-test shard that reports patch coverage. Claude-Session: https://claude.ai/code/session_01Ttd5Q9ZhRPB4ch5guos3rj --- litellm/router_strategy/lowest_latency.py | 54 +++---- .../test_lowest_latency_routing.py | 136 ------------------ .../router_strategy/test_lowest_latency.py | 127 ++++++++++++++++ 3 files changed, 149 insertions(+), 168 deletions(-) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index bb6877a032b..805d4ff9080 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -32,6 +32,12 @@ def _average_latency(samples: Sequence[float]) -> float: return sum(samples) / len(samples) +def _ttft_seconds(elapsed: timedelta | float) -> float: + if isinstance(elapsed, timedelta): + return elapsed.total_seconds() + return float(elapsed) + + class LowestLatencyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -86,14 +92,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms - time_to_first_token: float | None = None total_tokens = 0 if isinstance(response_obj, ModelResponse): @@ -111,13 +116,6 @@ class LowestLatencyLoggingHandler(CustomLogger): else: final_value = response_seconds - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) - # ------------ # Update usage # ------------ @@ -138,14 +136,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -252,7 +250,7 @@ class LowestLatencyLoggingHandler(CustomLogger): {model_group}_map: { id: { "latency": [..] - "time_to_first_token": [..] + "time_to_first_token_seconds": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -273,14 +271,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms total_tokens = 0 - time_to_first_token: float | None = None if isinstance(response_obj, ModelResponse): _usage: Final = getattr(response_obj, "usage", None) @@ -296,13 +293,6 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value = float(normalized_value) else: final_value = response_seconds - - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) # ------------ # Update usage # ------------ @@ -328,14 +318,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -433,7 +423,7 @@ class LowestLatencyLoggingHandler(CustomLogger): or float("inf") ) item_latency = item_map.get("latency", []) - item_ttft_latency = item_map.get("time_to_first_token", []) + item_ttft_latency = item_map.get("time_to_first_token_seconds", []) item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) diff --git a/tests/local_testing/test_lowest_latency_routing.py b/tests/local_testing/test_lowest_latency_routing.py index 598b1dbcaf9..aba4500199a 100644 --- a/tests/local_testing/test_lowest_latency_routing.py +++ b/tests/local_testing/test_lowest_latency_routing.py @@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async(): ), f"Oldest latency {oldest_latency} should have been discarded" -def test_ttft_list_trimming_discards_oldest_entry(): - """ - The time_to_first_token list trims the oldest entry when full, matching - the behavior of the latency list. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - # TTFT is only recorded when response_obj is a ModelResponse. - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - lowest_latency_logger.log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = test_cache.get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" - - @pytest.mark.asyncio async def test_timeout_penalty_discards_oldest_entry(): """ @@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims(): assert ( abs(latency_list[i] - expected) < tolerance ), f"At index {i}, expected ~{expected}, got {latency_list[i]}" - - -@pytest.mark.asyncio -async def test_ttft_list_trimming_discards_oldest_entry_async(): - """ - Async counterpart: the time_to_first_token list trims the oldest entry - when full. Exercises the async_log_success_event TTFT path, which only - runs when response_obj is a ModelResponse and the call is marked as - streaming with a completion_start_time. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - await lowest_latency_logger.async_log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = await test_cache.async_get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index eb02459be68..812d7bbff32 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -165,6 +165,133 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): json.dumps({"latency": latencies}) +MODEL_GROUP = "gpt-4o-mini" +FAST_TTFT_ID = "fast-ttft-short-output" +SLOW_TTFT_ID = "slow-ttft-long-output" +STREAMING_DEPLOYMENTS = [ + {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}}, + {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}}, +] + + +def _streaming_kwargs(deployment_id: str, start_time: datetime, ttft_seconds: float): + return { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": deployment_id}, + }, + "stream": True, + "completion_start_time": start_time + timedelta(seconds=ttft_seconds), + } + + +def _recorded_ttft(cache: DualCache, deployment_id: str): + cached = cache.get_cache(key=f"{MODEL_GROUP}_map") or {} + return cached.get(deployment_id, {}).get("time_to_first_token_seconds", []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_streaming_ttft_ranking_ignores_completion_length(sync_mode: bool): + """Deployment A: TTFT 1s, 50 completion tokens. Deployment B: TTFT 3s, 500 + completion tokens. Dividing TTFT by completion tokens made B look faster + (3/500 = 0.006 beats 1/50 = 0.02); actual TTFT must win.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + start_time = datetime(2026, 1, 1, 12, 0, 0) + end_time = start_time + timedelta(seconds=10) + + samples = ( + (FAST_TTFT_ID, 1.0, 50), + (SLOW_TTFT_ID, 3.0, 500), + ) + for deployment_id, ttft, completion_tokens in samples: + kwargs = _streaming_kwargs(deployment_id, start_time, ttft) + response_obj = _chat_response(completion_tokens=completion_tokens) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(1.0)] + assert _recorded_ttft(cache, SLOW_TTFT_ID) == [pytest.approx(3.0)] + + request_kwargs = {"stream": True, "metadata": {}} + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_ttft_window_keeps_newest_samples_when_full(sync_mode: bool): + """Float timestamps, as the SDK passes them. Once max_latency_list_size + samples exist the oldest TTFT is dropped so the window slides.""" + max_size = 3 + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": max_size}) + start_time = 1_700_000_000.0 + ttfts = (0.1, 0.2, 0.3, 0.4) + + for ttft in ttfts: + kwargs = { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": FAST_TTFT_ID}, + }, + "stream": True, + "completion_start_time": start_time + ttft, + } + response_obj = _chat_response(completion_tokens=1) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(ttft) for ttft in ttfts[-max_size:]] + + +@pytest.mark.asyncio +async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_workers(): + """Workers on the previous release share the Redis map and keep writing + seconds-per-token under the old "time_to_first_token" key during a rolling + deploy. Those samples favor SLOW; routing must only read the seconds key.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token": [0.02], "time_to_first_token_seconds": [1.0]}, + SLOW_TTFT_ID: {"time_to_first_token": [0.006], "time_to_first_token_seconds": [3.0]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", From 978f93ce9cf3e2748494a1c8b67b48488930b06b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:05:32 -0700 Subject: [PATCH 11/16] fix(tool_permission): log expected skip and deny events below WARNING (#40208) The no-tools skip path in async_pre_call_hook now logs at DEBUG and the denied-by-rule messages in async_pre_call_hook and _evaluate_tool_calls log at INFO. The malformed tool arguments warning is unchanged. Adds regression tests pinning each level Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/tool_permission.py | 6 +- .../guardrail_hooks/test_tool_permission.py | 75 +++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index a8b33109900..e20f0b320b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail): for _tool_call, is_allowed, _rule_id, message in checked: if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=message, blocked_content=True @@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail): new_tools: Final = self._collect_request_tools(data) if not new_tools: - verbose_proxy_logger.warning( + verbose_proxy_logger.debug( "Tool Permission Guardrail: not running guardrail. No tools or functions in data" ) return data @@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail): is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise HTTPException( status_code=400, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 0dbd4591ac9..87a1b84acc5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ import json +import logging import re from unittest.mock import patch @@ -520,6 +521,80 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_without_tools_logs_skip_at_debug(self, caplog): + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + result = await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + assert result is data + skip_levels = [r.levelno for r in caplog.records if "No tools or functions in data" in r.getMessage()] + assert skip_levels == [logging.DEBUG] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_pre_call_hook_denied_tool_logs_at_info(self, caplog): + data = {"tools": [{"type": "function", "function": {"name": "Read"}}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException): + await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_post_call_success_hook_denied_tool_logs_at_info(self, caplog): + tool_call = {"function": {"name": "Read", "arguments": "{}"}, "type": "function"} + response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})]) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data={"guardrails": ["test-tool-permission"]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + def test_parse_tool_call_arguments_malformed_json_logs_warning(self, caplog): + tool_call = ChatCompletionMessageToolCall(function={"name": "Bash", "arguments": "{not json"}, id="call_1") + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + parsed, error = self.guardrail._parse_tool_call_arguments(tool_call) + + assert parsed is None + assert error == "arguments could not be parsed" + warning_messages = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert len(warning_messages) == 1 + assert warning_messages[0].startswith("Tool Permission Guardrail: Failed to decode arguments for tool Bash") + @pytest.mark.asyncio async def test_async_pre_call_hook_blocks_legacy_functions(self): data = { From 99824533ff7a99c07e0c768f61fcd5b5204a7794 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:06:53 -0700 Subject: [PATCH 12/16] fix(proxy): kill the whole prisma process group when a boot migration command times out (#39509) * fix(proxy): kill the whole prisma process group when a boot migration command times out * fix(proxy): keep boot alive when the installed proxy extras has no prisma runner --- litellm/proxy/db/check_migration.py | 39 ++++++--- litellm/proxy/db/prisma_client.py | 22 ++++-- tests/test_litellm/proxy/db/conftest.py | 79 +++++++++++++++++++ .../proxy/db/test_check_migration.py | 32 ++++++++ .../proxy/db/test_prisma_client.py | 59 +++++++++----- 5 files changed, 197 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index b7a07d4eeea..6e2a06c96e1 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]: def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: """Checks for differences between current database and Prisma schema. + + Never raises: a diff that cannot be produced, because the runner is missing, + because the command failed, or because it outlived its budget, is reported as + "no diff" so boot continues. + Returns: A tuple containing: - A boolean indicating if differences were found (True) or not (False). - - A string with the diff output or error message. - Raises: - subprocess.CalledProcessError: If the Prisma command fails. - Exception: For any other errors during execution. + - The SQL commands that would close the diff, empty when there is none. """ - verbose_logger.debug("Checking for Prisma schema diff...") try: - result: Final = subprocess.run( + from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + ) + return False, [] + + verbose_logger.debug("Checking for Prisma schema diff...") + timeout: Final = prisma_command_timeout() + try: + result: Final = run_prisma( [ "prisma", "migrate", @@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: "./schema.prisma", "--script", ], - capture_output=True, - text=True, - check=True, + timeout=timeout, + env=os.environ.copy(), ) - # return True, "Migration diff generated successfully." sql_commands: Final = extract_sql_commands(result.stdout) if sql_commands: @@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: return True, sql_commands else: return False, [] + except subprocess.TimeoutExpired: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Timed out after {timeout}s generating the migration diff. " + f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." + ) + return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" print(error_message) # noqa: T201 diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 2190ae55fd2..9ea2432f2f5 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -937,9 +937,17 @@ class PrismaManager: use_v2_resolver=use_v2_resolver, ) else: + try: + from litellm_proxy_extras.prisma_toolchain import ( + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e) + return False + PrismaManager._raise_if_partitioned_spend_logs() - # Use prisma db push with increased timeout - subprocess.run( + run_prisma( [ "prisma", "db", @@ -947,13 +955,15 @@ class PrismaManager: "--accept-data-loss", "--skip-generate", ], - timeout=60, - check=True, + timeout=prisma_command_timeout(), + env=os.environ.copy(), + stdout=None, + stderr=None, ) PrismaManager._apply_replica_identity_full_if_requested() return True - except subprocess.TimeoutExpired: - verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1) + except subprocess.TimeoutExpired as e: + verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index 6f67b91ac1d..03d7ea81257 100644 --- a/tests/test_litellm/proxy/db/conftest.py +++ b/tests/test_litellm/proxy/db/conftest.py @@ -1,5 +1,11 @@ +import json import os +import signal +import sys +import time from collections.abc import Generator +from dataclasses import dataclass +from pathlib import Path from typing import Optional import pytest @@ -75,3 +81,76 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]: def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DATABASE_URL", "about-to-be-unset") monkeypatch.delenv("DATABASE_URL") + + +FAKE_PRISMA_CLI = """#!{python} +import json +import os +import pathlib +import subprocess +import sys +import time + +calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"]) +earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else [] +with calls_file.open("a") as log: + print(json.dumps(sys.argv[1:]), file=log) +if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"): + grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid)) + time.sleep(600) +sys.exit(0) +""" + + +@dataclass(frozen=True, slots=True) +class FakePrismaCli: + """A stand-in `prisma` on PATH, recording every invocation. + + With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree + of its own, the way the real CLI wraps Node around a Rust schema engine, so a + timeout that kills only the direct child leaves the rest of that tree running. + """ + + calls_file: Path + grandchild_pidfile: Path + + @property + def calls(self) -> list[list[str]]: + if not self.calls_file.exists(): + return [] + return [json.loads(line) for line in self.calls_file.read_text().splitlines()] + + def grandchild_is_gone(self, within_seconds: float) -> bool: + deadline = time.monotonic() + within_seconds + while time.monotonic() < deadline: + try: + os.kill(int(self.grandchild_pidfile.read_text()), 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +@pytest.fixture +def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]: + bin_dir = tmp_path / "fakebin" + bin_dir.mkdir() + script = bin_dir / "prisma" + script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable)) + script.chmod(0o755) + cli = FakePrismaCli( + calls_file=tmp_path / "calls.jsonl", + grandchild_pidfile=tmp_path / "grandchild.pid", + ) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") + monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file)) + monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile)) + monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1") + monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False) + yield cli + if cli.grandchild_pidfile.exists(): + try: + os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL) + except ProcessLookupError: + pass diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py index 9e2f6a1089c..74a6841cb23 100644 --- a/tests/test_litellm/proxy/db/test_check_migration.py +++ b/tests/test_litellm/proxy/db/test_check_migration.py @@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker): check_migration.verbose_logger.exception.assert_called_once() actual_message = check_migration.verbose_logger.exception.call_args[0][0] assert "prisma schema out of sync with db" in actual_message + + +@pytest.mark.timeout(30) +def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch): + """ + `prisma migrate diff` ran unbounded, so a database that never answers hung boot + before uvicorn ever started, and interrupting the proxy orphaned the schema engine. + """ + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + assert fake_prisma_cli.calls == [ + ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x", + "--to-schema-datamodel", "./schema.prisma", "--script"] + ] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch): + """ + Boot calls this helper directly, so an ImportError here takes the proxy down before + uvicorn starts. An install without the runner must lose the diagnostic, not the proxy. + """ + import sys + + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index f0983d6bf62..7f43e483557 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -10,7 +10,7 @@ from fastapi.testclient import TestClient -from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema +from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema @pytest.fixture(autouse=True) @@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client( mock_new_prisma.connect.assert_awaited_once() -def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): +DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"] + + +def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url): """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call into the opt-in REPLICA IDENTITY FULL step.""" from litellm.proxy.db.prisma_client import PrismaManager @@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): staticmethod(lambda: applied.append(True)), ) - with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] assert applied == [True] -def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): +def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the primary key back to ("request_id"), which Postgres rejects; the guard must fail fast with guidance instead of running the push.""" @@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - with pytest.raises(RuntimeError) as err: - PrismaManager.setup_database(use_migrate=False) + with pytest.raises(RuntimeError) as err: + PrismaManager.setup_database(use_migrate=False) assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR - mock_run.assert_not_called() + assert fake_prisma_cli.calls == [] -def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch): +def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): from litellm.proxy.db.prisma_client import PrismaManager from litellm_proxy_extras.utils import ProxyExtrasDBManager monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] def _entra_jwt(expires_in_seconds: int) -> str: @@ -377,3 +373,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch with pytest.raises(RuntimeError, match="DATABASE_HOST"): wrapper.get_rds_iam_token() + + +@pytest.mark.timeout(45) +def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch): + """ + A timed-out `db push` used to leave Node and the schema engine writing the schema, + so the next attempt pushed into a database the abandoned one was still mutating. + """ + monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False) + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert PrismaManager.setup_database(use_migrate=False) is True + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot( + fake_prisma_cli, unset_database_url, monkeypatch +): + """ + An ImportError out of setup_database escapes the caller's RuntimeError handler and + kills boot, bypassing the operator's enforce_prisma_migration_check choice. + """ + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert PrismaManager.setup_database(use_migrate=False) is False + assert fake_prisma_cli.calls == [] From a85c3152ca59b5b56d1168e2a384e29a4fb0b861 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:08:24 -0700 Subject: [PATCH 13/16] fix(proxy): load db credentials in the model reconcile so a worker never serves a model before its credential (#39876) * fix(proxy): load db credentials inside the model reconcile so a worker never serves a model before its credential * fix(proxy): load db credentials in the model read-through so a request miss never adds a model before its credential * fix(proxy): read credentials from the writer db before the router update and look a credential up once * test(proxy): assert the credential is loaded when db models reach the router instead of the call order --- .../litellm_core_utils/credential_accessor.py | 15 ++- .../common_utils/registry_read_through.py | 1 + litellm/proxy/proxy_server.py | 31 +---- litellm/utils.py | 18 ++- .../common_utils/test_config_sync_pubsub.py | 7 +- .../test_registry_read_through.py | 45 +++++++ .../test_scheduled_job_stagger.py | 2 +- .../proxy/proxy_server/test_proxy_config.py | 123 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 68 +++++----- tests/test_litellm/test_utils.py | 44 +++++++ 10 files changed, 278 insertions(+), 76 deletions(-) diff --git a/litellm/litellm_core_utils/credential_accessor.py b/litellm/litellm_core_utils/credential_accessor.py index 7071750b970..7b4e8240b69 100644 --- a/litellm/litellm_core_utils/credential_accessor.py +++ b/litellm/litellm_core_utils/credential_accessor.py @@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem class CredentialAccessor: + @staticmethod + def find_credential(credential_name: str) -> CredentialItem | None: + return next( + (credential for credential in litellm.credential_list if credential.credential_name == credential_name), + None, + ) + @staticmethod def get_credential_values(credential_name: str) -> dict: """Safe accessor for credentials.""" - if not litellm.credential_list: - return {} - for credential in litellm.credential_list: - if credential.credential_name == credential_name: - return credential.credential_values.copy() - return {} + credential: Final = CredentialAccessor.find_credential(credential_name) + return {} if credential is None else credential.credential_values.copy() @staticmethod def upsert_credentials(credentials: list[CredentialItem]): diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 460b348e188..63da5d15207 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool: ) return proxy_server.llm_router is not None async with proxy_server.MODEL_RECONCILE_LOCK: + await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client) proxy_server.proxy_config._add_deployment(db_models=rows) proxy_server.llm_model_list = router.get_model_list() return True diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 32b6b841af7..1d6baf5fded 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7109,11 +7109,10 @@ class ProxyConfig: ], ) - # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) - if self._should_load_db_object(object_type="models"): - new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) - - # update llm router + load_models: Final = self._should_load_db_object(object_type="models") + new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None + await self.get_credentials(prisma_client=prisma_client) + if load_models: still_desired_ids = await self._update_llm_router( new_models=new_models, proxy_logging_obj=proxy_logging_obj ) @@ -7153,12 +7152,9 @@ class ProxyConfig: async def _resync_config_from_db() -> None: await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) - async def _resync_credentials_from_db() -> None: - await self.get_credentials(prisma_client=prisma_client) - subscriber: Final = ConfigSyncSubscriber( redis_cache=redis_cache, - resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + resync_callbacks=(_resync_config_from_db,), ) self.config_sync_subscriber = subscriber subscriber.start() @@ -8013,7 +8009,7 @@ class ProxyConfig: async def get_credentials(self, prisma_client: PrismaClient): try: - credentials = await CredentialsRepository(prisma_client).find_all() + credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] await self.delete_credentials(credentials) # delete credentials that are not in the all-up list CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list @@ -9597,19 +9593,6 @@ class ProxyStartupEvent: ) if store_model_in_db is True: - ### GET STORED CREDENTIALS ### - scheduler.add_job( - proxy_config.get_credentials, - "interval", - seconds=config_reload_interval_seconds, - # REMOVED jitter parameter - major cause of memory leak - args=[prisma_client], - id="get_credentials_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - await proxy_config.get_credentials(prisma_client=prisma_client) - # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum # Frequent polling was causing excessive memory allocations scheduler.add_job( @@ -9623,7 +9606,7 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - # this will load all existing models on proxy startup + # this will load all existing credentials and models on proxy startup await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) proxy_config.start_config_sync_subscriber( diff --git a/litellm/utils.py b/litellm/utils.py index d0e11bc9551..141ec323776 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -672,11 +672,19 @@ def load_credentials_from_list(kwargs: dict): CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor") credential_name: Final = kwargs.get("litellm_credential_name") - if credential_name and litellm.credential_list: - credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name) - for key, value in credential_accessor.items(): - if key not in kwargs: - kwargs[key] = value + if not credential_name: + return + credential: Final = CredentialAccessor.find_credential(credential_name) + if credential is None: + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + credential_name, + len(litellm.credential_list), + ) + return + for key, value in credential.credential_values.items(): + if key not in kwargs: + kwargs[key] = value def get_dynamic_callbacks( diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 3bb3f75b9b8..2a0ebf9492a 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -824,7 +824,7 @@ class _StopFailingSubscriber(ConfigSyncSubscriber): raise RuntimeError("stop failed") -async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None: +async def test_proxy_config_subscriber_resyncs_deployments_only() -> None: from litellm.proxy.proxy_server import ProxyConfig cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) @@ -852,10 +852,7 @@ async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> await callback() await config.stop_config_sync_subscriber() - assert calls == [ - ("add_deployment", prisma_client, proxy_logging_obj), - ("get_credentials", prisma_client, None), - ] + assert calls == [("add_deployment", prisma_client, proxy_logging_obj)] assert config.config_sync_subscriber is None assert subscriber._task is None diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index f0fbdea4e85..6f7c20166c5 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -393,6 +393,51 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc assert not proxy_server.MODEL_RECONCILE_LOCK.locked() +@pytest.mark.asyncio +async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import AsyncMock, MagicMock + + import litellm + import litellm.proxy.proxy_server as proxy_server + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + from litellm.types.utils import CredentialItem + + rows: Final = [MagicMock()] + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) + router: Final = MagicMock() + router.get_model_list.return_value = [] + installed: Final = MagicMock() + + async def load_credentials_from_db(prisma_client: object) -> None: + CredentialAccessor.upsert_credentials( + [ + CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db"}, + credential_info={}, + ) + ] + ) + + def install_models(db_models: object) -> None: + installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models) + + assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True + installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"}) + + @pytest.mark.asyncio async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index ca4d62737b6..54e0aa74a25 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py @@ -23,7 +23,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import ( ) OPERATOR_CRON_JOB_ID = "spend_log_cleanup_job" -SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "get_credentials_job", "add_deployment_job") +SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "add_deployment_job") async def _noop() -> None: ... diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 770cec1834e..ae4ec4086bb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2939,6 +2939,129 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing") +def _stub_add_deployment_collaborators( + monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock +) -> None: + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.get_model_list = MagicMock(return_value=[]) + + async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]: + return {} + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(pc, "get_config", fake_get_config) + monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock()) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "master_key", "sk-master") + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "proxy_config", pc) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + + +def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + return { + "credential_name": credential_name, + "credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock: + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row(credential_name, api_key)] + ) + return fake_prisma + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + from litellm.utils import load_credentials_from_list + + pc = ProxyConfig() + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {}) + installed = MagicMock() + + async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]: + fake_prisma.db.litellm_credentialstable.find_many.return_value = [ + _encrypted_credential_row("openai-cred", "sk-from-db") + ] + return [MagicMock()] + + async def install_models(new_models: object, proxy_logging_obj: object) -> None: + installed(credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands) + monkeypatch.setattr(pc, "_update_llm_router", install_models) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + installed.assert_called_once_with(credential={"api_key": "sk-from-db"}) + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + request_kwargs = {"litellm_credential_name": "openai-cred"} + load_credentials_from_list(request_kwargs) + assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + + pc = ProxyConfig() + fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db") + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]}) + models_fetch = AsyncMock(return_value=[]) + monkeypatch.setattr(pc, "_get_models_from_db", models_fetch) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + models_fetch.assert_not_awaited() + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + pc = ProxyConfig() + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")] + ) + reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + fake_prisma = MagicMock() + fake_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + + await pc.get_credentials(prisma_client=fake_prisma) + + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"} + reader_inner.litellm_credentialstable.find_many.assert_not_awaited() + + # --------------------------------------------------------------------------- # ProxyConfig._add_general_settings_from_db_config # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b7bb58378d4..937e4f15741 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -822,16 +822,16 @@ def _mock_scheduled_proxy_config() -> MagicMock: @pytest.mark.asyncio -async def test_initialize_scheduled_jobs_credentials(monkeypatch): - """ - Test that get_credentials is only called when store_model_in_db is True - """ +async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_deployment( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging - # Mock dependencies mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() @@ -841,25 +841,6 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch("litellm.proxy.proxy_server.store_model_in_db", False), - ): # set store_model_in_db to False - # Test when store_model_in_db is False - await ProxyStartupEvent.initialize_scheduled_background_jobs( - general_settings={}, - prisma_client=mock_prisma_client, - proxy_budget_rescheduler_min_time=1, - proxy_budget_rescheduler_max_time=2, - proxy_batch_write_at=5, - proxy_logging_obj=mock_proxy_logging, - ) - - # Verify get_credentials was not called - mock_proxy_config.get_credentials.assert_not_called() - - # Now test with store_model_in_db = True - with ( - patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), - patch("litellm.proxy.proxy_server.store_model_in_db", True), - patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True), ): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, @@ -870,12 +851,31 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): proxy_logging_obj=mock_proxy_logging, ) - # Verify get_credentials was called both directly and scheduled - assert mock_proxy_config.get_credentials.call_count == 1 # Direct call + mock_proxy_config.get_credentials.assert_not_called() + mock_proxy_config.add_deployment.assert_not_called() - # Verify a scheduled job was added for get_credentials - mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls] - assert len(mock_scheduler_calls) > 0 + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + assert scheduler.get_job("get_credentials_job") is None + assert scheduler.get_job("add_deployment_job") is not None + mock_proxy_config.get_credentials.assert_not_called() + assert mock_proxy_config.add_deployment.call_count == 1 + finally: + scheduler.shutdown(wait=False) @pytest.mark.asyncio @@ -924,7 +924,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ - The DB config-reload jobs (add_deployment, get_credentials) that keep multi-pod + The DB config-reload job (add_deployment) that keeps multi-pod deployments in sync must be scheduled at the configured proxy_config_reload_interval_seconds, not a hardcoded value. """ @@ -967,7 +967,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval( if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == configured_interval - assert scheduled_seconds["get_credentials_job"] == configured_interval + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -1011,7 +1011,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == 30 - assert scheduled_seconds["get_credentials_job"] == 30 + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -7446,10 +7446,8 @@ async def test_store_model_in_db_db_override_when_config_false(): # store_model_in_db should now be True (overridden by DB) assert ps.store_model_in_db is True - # add_deployment and get_credentials should have been called - # since store_model_in_db is now True assert mock_proxy_config.add_deployment.call_count == 1 - assert mock_proxy_config.get_credentials.call_count == 1 + mock_proxy_config.get_credentials.assert_not_called() @pytest.mark.asyncio diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8a56a84ade7..4ce68b71079 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -6092,3 +6092,47 @@ class TestFinalOptionalParamsLineRedaction: assert "'max_tokens': 17" in printed assert "'temperature': 0.25" in printed + + +def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()] + + +def test_load_credentials_from_list_warns_when_the_named_credential_is_not_loaded( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.utils import load_credentials_from_list + + monkeypatch.setattr(litellm, "credential_list", []) + request_kwargs = {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + assert _credential_warnings(caplog) == [ + "litellm_credential_name=openai-cred matched none of the 0 loaded credentials; the request runs without it" + ] + + +def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.types.utils import CredentialItem + from litellm.utils import load_credentials_from_list + + loaded = CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db", "api_base": "https://credential.example"}, + credential_info={}, + ) + monkeypatch.setattr(litellm, "credential_list", [loaded]) + request_kwargs = {"litellm_credential_name": "openai-cred", "api_base": "https://request.example"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == { + "litellm_credential_name": "openai-cred", + "api_base": "https://request.example", + "api_key": "sk-from-db", + } + assert _credential_warnings(caplog) == [] From f769aa4675de225a04fdc055c092318d94612502 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 10:11:20 -0700 Subject: [PATCH 14/16] fix(router): give cooldowns their own cache so siblings see a bench in ~1s (#40025) Cooldown entries rode the router-wide DualCache, which re-reads a key that is missing from memory at most once every 10s. A deployment benched on one replica therefore kept taking traffic on its siblings for up to 10 seconds, and the same shared in-memory tier could evict a live cooldown once 200 unrelated router keys crowded it out, which sent even the benching replica back to the dead deployment. CooldownCache now owns a DualCache over the router's Redis with a 1s read interval and an in-memory tier that only holds cooldown keys. Redis is attached lazily because the router builds the cooldown cache before it wires Redis up. --- litellm/constants.py | 3 + litellm/router_utils/cooldown_cache.py | 48 +++-- .../test_router_cooldown_per_deployment.py | 16 +- .../router_utils/test_cooldown_cache.py | 184 ++++++++++++++++-- 4 files changed, 212 insertions(+), 39 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index defc9337e9b..8ef9523a60a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -73,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) +DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float( + os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1") +) DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 9e7f457f631..ef29f7d8fd3 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from litellm import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker if TYPE_CHECKING: @@ -36,10 +37,19 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0 class CooldownCache: - def __init__(self, cache: DualCache, default_cooldown_time: float): + def __init__( + self, + cache: DualCache, + default_cooldown_time: float, + redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, + ): self.cache = cache self.default_cooldown_time = default_cooldown_time self.in_memory_cache = InMemoryCache() + self._cooldown_store = DualCache( + in_memory_cache=self.in_memory_cache, + default_redis_batch_cache_expiry=redis_read_interval_seconds, + ) # Initialize the masker with custom settings for exception strings self.exception_masker = SensitiveDataMasker( visible_prefix=50, # Show first 50 characters @@ -48,6 +58,21 @@ class CooldownCache: mask_short_values=False, # Truncate long messages only; keep short ones readable ) + @property + def cooldown_store(self) -> DualCache: + """ + The cache cooldown entries live in, with the router's Redis attached on first use. + + It is kept separate from the router-wide cache so that a key missing from memory is + re-read from Redis every `redis_read_interval_seconds` rather than on the router + cache's much longer batch interval, which is what lets a sibling replica see a + cooldown another replica wrote, and so that unrelated router keys cannot evict a + cooldown from the in-memory tier before it expires. Redis is attached lazily because + the router builds its cooldown cache before it wires up the shared Redis client. + """ + self._cooldown_store.attach_redis_cache(self.cache.redis_cache) + return self._cooldown_store + def _common_add_cooldown_logic( self, model_id: str, original_exception, exception_status, cooldown_time: float ) -> tuple[str, CooldownCacheValue]: @@ -93,7 +118,7 @@ class CooldownCache: ) # Set the cache with a TTL equal to the cooldown time - self.cache.set_cache( + self.cooldown_store.set_cache( value=cooldown_data, key=cooldown_key, ttl=_cooldown_time, @@ -122,13 +147,13 @@ class CooldownCache: cooldown_cache_value: Final = CooldownCacheValue(**result) # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time if remaining <= 0: - self.cache.in_memory_cache.delete_cache(key) + self.in_memory_cache.delete_cache(key) return None - current_expiry: Final = self.cache.in_memory_cache.ttl_dict.get(key) + current_expiry: Final = self.in_memory_cache.ttl_dict.get(key) if current_expiry is not None and current_expiry > current_time + remaining + 5: corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS) - self.cache.in_memory_cache.delete_cache(key) - self.cache.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) + self.in_memory_cache.delete_cache(key) + self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) return cooldown_cache_value async def async_get_active_cooldowns( @@ -137,12 +162,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] - # Retrieve the values for the keys using mget - ## more likely to be none if no models ratelimited. So just check redis every 1s - ## each redis call adds ~100ms latency. - - ## check in memory cache first - results: Final = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) + results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = [] if results is None or all(v is None for v in results): @@ -164,7 +184,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] active_cooldowns: Final = [] current_time: Final = time.time() @@ -184,7 +204,7 @@ class CooldownCache: keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] min_cooldown_time: float | None = None # Process the results diff --git a/tests/router_unit_tests/test_router_cooldown_per_deployment.py b/tests/router_unit_tests/test_router_cooldown_per_deployment.py index b8ae8a8c013..228dfed38b9 100644 --- a/tests/router_unit_tests/test_router_cooldown_per_deployment.py +++ b/tests/router_unit_tests/test_router_cooldown_per_deployment.py @@ -246,12 +246,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -267,7 +267,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -290,14 +290,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -318,12 +318,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None class TestFallbackDeploymentCooldown: diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/test_litellm/router_utils/test_cooldown_cache.py index 68e9aeaa4fc..6f90fa8465f 100644 --- a/tests/test_litellm/router_utils/test_cooldown_cache.py +++ b/tests/test_litellm/router_utils/test_cooldown_cache.py @@ -268,12 +268,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -289,7 +289,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -312,14 +312,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -340,12 +340,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None @pytest.mark.asyncio async def test_async_active_entry_is_returned(self): @@ -363,7 +363,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -389,18 +389,18 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:expired-dep:cooldown" entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is None - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None def test_active_entry_within_window_returns_value(self): cc = self._make_cooldown_cache() key = "deployment:active-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) @@ -412,12 +412,12 @@ class TestCorrectedActiveCooldown: key = "deployment:backfilled-dep:cooldown" remaining = 30.0 entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is not None - corrected_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + corrected_expiry = cc.in_memory_cache.ttl_dict.get(key) assert corrected_expiry is not None assert corrected_expiry - time.time() <= 60.0 @@ -425,10 +425,160 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:normal-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) - original_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) + original_expiry = cc.in_memory_cache.ttl_dict.get(key) cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry == original_expiry + + +class SharedRedisDouble: + """ + In-process stand-in for RedisCache, shared by several DualCache instances so that + tests can model two proxy replicas talking to one Redis. + """ + + def __init__(self) -> None: + self.store: dict = {} # mutable-ok: stands in for Redis' own mutable keyspace + + def set_cache(self, key, value, **kwargs): + self.store[key] = value + + async def async_set_cache(self, key, value, **kwargs): + self.store[key] = value + + def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + +class TestCooldownPropagationBetweenReplicas: + """ + A cooldown written by one replica has to reach its siblings quickly. The router's own + DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads + get their own cache with a much shorter Redis read interval. + """ + + def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache: + router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis) + if read_interval is None: + return CooldownCache(cache=router_cache, default_cooldown_time=60.0) + return CooldownCache( + cache=router_cache, + default_cooldown_time=60.0, + redis_read_interval_seconds=read_interval, + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "shared-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "sibling replica must pick up a cooldown written by another replica within the read interval" + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_default_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis) + replica_b = self._make_replica(redis) + model_id = "default-interval-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(1.2) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "the shipped default read interval must let a sibling replica see a cooldown about a second later" + ) + + def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "sync-shared-deployment" + + assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active] + + @pytest.mark.asyncio + async def test_redis_attached_after_construction_is_still_used(self): + redis = SharedRedisDouble() + router_cache = DualCache(in_memory_cache=InMemoryCache()) + writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25) + router_cache.attach_redis_cache(redis) + reader = self._make_replica(redis, read_interval=0.25) + model_id = "late-redis-deployment" + + writer.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "a router that wires Redis after building its cooldown cache must still publish cooldowns to it" + ) + + +class TestCooldownSurvivesUnrelatedCacheTraffic: + @pytest.mark.asyncio + async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self): + router_cache = DualCache(in_memory_cache=InMemoryCache()) + cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0) + model_id = "busy-router-deployment" + + cc.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=30.0, + ) + + for i in range(400): + router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i}) + + active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "unrelated router cache traffic must not evict a cooldown that is still running" + ) From ee1a6407cb5b901660e62d757f83cdd6a99e5a34 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:45:02 -0700 Subject: [PATCH 15/16] fix(guardrails): keep guardrail telemetry when a policy pipeline blocks or modifies the response (#40211) * fix(guardrails): keep guardrail telemetry when a policy pipeline blocks or modifies the response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): count every raw-snapshot guardrail evaluation and type the telemetry carry helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(policy_engine): type the recording guardrail hooks and telemetry test parameters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/pipeline_executor.py | 104 +++++++++---- .../policy_engine/test_pipeline_executor.py | 141 ++++++++++++++++++ 2 files changed, 216 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4be0f556ed7..1bcda38657b 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal import litellm @@ -16,7 +16,11 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, + independent_snapshot, +) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -25,6 +29,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStep, PipelineStepResult, ) +from litellm.types.utils import StandardLoggingGuardrailInformation try: from fastapi.exceptions import HTTPException @@ -118,6 +123,7 @@ class PipelineExecutor: return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="block", step_results=step_results, @@ -126,6 +132,7 @@ class PipelineExecutor: ) if action == "modify_response": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="modify_response", step_results=step_results, @@ -168,34 +175,33 @@ class PipelineExecutor: verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) return ("error", None, f"Guardrail '{step.guardrail}' not found", None) + # Inject guardrail name into metadata so should_run_guardrail() allows it + if "metadata" not in data: + data["metadata"] = {} + data["metadata"]["guardrails"] = [step.guardrail] + + # A scan_raw_request step evaluates the pristine pre-pipeline + # snapshot instead of `data` (which earlier pass_data steps in + # this same pipeline may have already rewritten), same reason + # the normal sequential/parallel guardrail loops do this. + scans_raw_request: Final = callback.scan_raw_request + hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data + independent_snapshot(raw_request_snapshot) + if scans_raw_request and raw_request_snapshot is not None + else data + ) + if hook_input is not data: + hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] + snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input)) + + # Use unified_guardrail path if callback implements apply_guardrail + target: CustomLogger = callback + use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + if use_unified: + hook_input["guardrail_to_apply"] = callback + target = UnifiedLLMGuardrails() + try: - # Inject guardrail name into metadata so should_run_guardrail() allows it - if "metadata" not in data: - data["metadata"] = {} - data["metadata"]["guardrails"] = [step.guardrail] - - # A scan_raw_request step evaluates the pristine pre-pipeline - # snapshot instead of `data` (which earlier pass_data steps in - # this same pipeline may have already rewritten), same reason - # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = callback.scan_raw_request - hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data - independent_snapshot(raw_request_snapshot) - if scans_raw_request and raw_request_snapshot is not None - else data - ) - if hook_input is not data: - hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] - - # Use unified_guardrail path if callback implements apply_guardrail - target: CustomLogger = callback - use_unified: Final = ( - "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks - ) - if use_unified: - hook_input["guardrail_to_apply"] = callback - target = UnifiedLLMGuardrails() - if mode == "pre_call": response = await target.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -233,6 +239,12 @@ class PipelineExecutor: else: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + finally: + if hook_input is not data: + _append_guardrail_information( + request_data=data, + entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:], + ) @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: @@ -283,6 +295,40 @@ def _restore_request_guardrails( return {**working_data, "metadata": stripped} # mutable-ok: request dict +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]: + bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source)) + recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + return recorded if isinstance(recorded, list) else [] + + +def _append_guardrail_information( + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data + entries: Sequence[StandardLoggingGuardrailInformation], +) -> None: + if not entries: + return + _, request_bucket = get_or_create_metadata_bucket(request_data) + existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY) + if isinstance(existing, list): + existing.extend(entries) + return + request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries) + + +def _carry_working_guardrail_information( + working_data: Mapping[str, object], + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data +) -> None: + recorded: Final = _recorded_guardrail_information(working_data) + existing: Final = _recorded_guardrail_information(request_data) + if recorded is existing: + return + _append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing]) + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 4fcb7d22588..e95d001bb12 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -4,20 +4,26 @@ Tests for the pipeline executor. Uses mock guardrails to validate pipeline execution without external services. """ +import copy +from typing import Literal from unittest.mock import MagicMock import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( CustomCodeGuardrail, ) from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) +from litellm.types.utils import CallTypesLiteral try: from fastapi.exceptions import HTTPException @@ -158,11 +164,146 @@ class ContentCheckGuardrail(CustomGuardrail): return None +class RecordingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + scan_raw_request=scan_raw_request, + ) + self.block = block + + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"detected": ["aws_access_key"]}, + request_data=data, + guardrail_status="guardrail_intervened" if self.block else "success", + ) + if self.block: + raise HTTPException(status_code=400, detail="Content policy violation") + return copy.deepcopy(data) + + # ───────────────────────────────────────────────────────────────────────────── # Tests # ───────────────────────────────────────────────────────────────────────────── +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +@pytest.mark.parametrize("scan_raw_request", [False, True]) +@pytest.mark.parametrize("on_fail", ["block", "modify_response"]) +async def test_terminal_block_carries_guardrail_information_to_request( + monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"] +): + """ + Spend logging and the Guardrails Monitor read standard_logging_guardrail_information + off the caller's request dict. A blocking step records it on the executor's + working copy (or the raw-request snapshot), so the terminal result must carry it + back onto the request or the block is never counted. + """ + guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request) + monkeypatch.setattr(litellm, "callbacks", [guard]) + data = { + "messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}], + "metadata": {"user_api_key_hash": "abc"}, + } + + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}}, + ) + + assert result.terminal_action == on_fail + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"] + assert recorded[0]["guardrail_status"] == "guardrail_intervened" + assert data["metadata"]["user_api_key_hash"] == "abc" + assert "guardrails" not in data["metadata"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch): + """A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step + that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two + dicts at once. Those must be carried back once while every step's own entry is kept.""" + first = RecordingGuardrail(guardrail_name="pii-scan", block=False) + second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "block" + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch): + """Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller, + while the entries the raw snapshot already held before the pipeline ran are not copied again.""" + guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False) + monkeypatch.setattr(litellm, "callbacks", [guard]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="raw-scan-policy", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"] + + @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio async def test_escalation_step1_fails_step2_blocks(monkeypatch): From 35d1d40a6784c222714e880c7d07729bc90018b2 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Tue, 8 Sep 2026 12:07:25 -0700 Subject: [PATCH 16/16] fix(ocr): run post-call logging hooks (#40154) --- litellm/integrations/custom_guardrail.py | 10 +-- litellm/llms/custom_httpx/llm_http_handler.py | 12 ++++ .../test_secret_detection.py | 2 +- .../integrations/test_custom_guardrail.py | 33 ++++++++++ .../custom_httpx/test_llm_http_handler.py | 64 +++++++++++++++++++ .../test_proxy_logging_hook_detection.py | 10 +-- 6 files changed, 121 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2d66a280663..e76d4b05bb2 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -773,7 +773,7 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def _deployment_pre_call_target(self) -> "CustomLogger": + def _deployment_hook_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self try: @@ -802,7 +802,7 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - target: Final = self._deployment_pre_call_target() + target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self result: Final = await target.async_pre_call_hook( @@ -845,7 +845,9 @@ class CustomGuardrail(CustomLogger): return None # CHECK IF GUARDRAIL REJECTS THE REQUEST - result: Final = await self.async_post_call_success_hook( + target: Final = self._deployment_hook_target() + hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data + result: Final = await target.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth( user_id=request_data.get("user_api_key_user_id"), team_id=request_data.get("user_api_key_team_id"), @@ -853,7 +855,7 @@ class CustomGuardrail(CustomLogger): api_key=request_data.get("user_api_key_hash"), request_route=request_data.get("user_api_key_request_route"), ), - data=request_data, + data=hook_request_data, response=response, ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..8b6d3a0ceb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + return self._transform_ocr_response( provider_config=provider_config, model=model, @@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + # Use async response transform for async operations return await provider_config.async_transform_ocr_response( model=model, diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index f46df5baadf..c1c3569c0e2 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks(): never sees ``data["prompt"]``.""" guardrail = _guardrail() assert guardrail.uses_apply_guardrail_interface() is True - assert guardrail._deployment_pre_call_target() is guardrail + assert guardrail._deployment_hook_target() is guardrail @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 885bd1d4d72..359635335b3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2610,3 +2610,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: ) assert result is replacement + + @pytest.mark.asyncio + async def test_apply_guardrail_interface_modifies_deployment_response(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + class ReplacingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + assert input_type == "response" + return {**inputs, "texts": ["filtered response"]} + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}]) + request_data = {"guardrails": ["test-guardrail"]} + + result = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is response + assert response.choices[0].message.content == "filtered response" + assert request_data == {"guardrails": ["test-guardrail"]} diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index e16855da8cb..98a5f4b2db5 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams @@ -37,6 +38,69 @@ from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, Trans _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "OCR output", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +def _ocr_sync_client() -> HTTPHandler: + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))) + return client + + +def _ocr_async_client() -> AsyncHTTPHandler: + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)) + ) + return client + + +def test_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = BaseLLMHTTPHandler().ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_sync_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + + +@pytest.mark.asyncio +async def test_async_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = await BaseLLMHTTPHandler().async_ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_async_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..22930a26974 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == [] -def test_deployment_pre_call_target_stays_native_when_opted_out(): +def test_deployment_hook_target_stays_native_when_opted_out(): """Model-level guardrails resolve their target here rather than through ProxyLogging.""" - assert _KeepsNativeHooks()._deployment_pre_call_target() is not None + assert _KeepsNativeHooks()._deployment_hook_target() is not None opted_out = _KeepsNativeHooks() - assert opted_out._deployment_pre_call_target() is opted_out - assert _AppliesGuardrail()._deployment_pre_call_target() is not None + assert opted_out._deployment_hook_target() is opted_out + assert _AppliesGuardrail()._deployment_hook_target() is not None routed = _AppliesGuardrail() - assert routed._deployment_pre_call_target() is not routed + assert routed._deployment_hook_target() is not routed @pytest.mark.asyncio