scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_responses_ws_spend_tracking

Browse source
This commit is contained in:
yucheng 2026-09-08 17:39:12 +00:00
commit ca3916b338
14 changed files with 1146 additions and 230 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.circleci/config.yml
View file

@ -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 \

130
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -8,14 +8,10 @@ import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
from typing import TYPE_CHECKING, Final, Optional
from litellm_proxy_extras import prisma_toolchain
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 (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
prisma_command_timeout,
prisma_migrate_deploy_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:
@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
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
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@ -624,7 +650,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:
@ -645,7 +671,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:
@ -719,6 +745,95 @@ 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. 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:
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(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")
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:
"""
@ -994,6 +1109,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

65
litellm/caching/redis_cache.py
View file

@ -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,29 @@ 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(
{"<lambda>", "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 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"
)
_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:
"""
@ -736,6 +756,43 @@ 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 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.
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
_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl
)
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()
@ -1241,6 +1298,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)

10
litellm/router.py
View file

@ -1248,7 +1248,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:
@ -4214,10 +4214,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),
messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally
specific_deployment=specific_deployment,
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs)

378
litellm/router_strategy/least_busy.py
View file

@ -1,17 +1,103 @@
#### 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])
_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | 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_counts(values: Sequence[float | None]) -> tuple[int, ...]:
return tuple(0 if value is None else int(value) for value in values)
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(
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 shared in-flight counts for %s, "
"falling back to this worker's own counts: %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
@ -20,195 +106,101 @@ 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, 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)
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is not None:
try:
shared: Final = _as_counts(redis_cache.batch_get_counts(list(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: 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)
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is not None:
try:
shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(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)
if ref is None:
return
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:
return
redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
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:
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:
return
await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
except Exception as e:
_warn_unwritable(key, e)

10
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 2956
"limit": 2918
},
"ANN002": {
"limit": 71
@ -9,10 +9,10 @@
"limit": 806
},
"ANN201": {
"limit": 1979
"limit": 1965
},
"ANN202": {
"limit": 831
"limit": 829
},
"ANN204": {
"limit": 683
@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
"limit": 2916
"limit": 2914
},
"C401": {
"limit": 8
@ -189,7 +189,7 @@
"limit": 0
},
"S110": {
"limit": 217
"limit": 207
},
"S112": {
"limit": 22

46
tests/local_testing/test_least_busy_routing.py
View file

@ -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}"

80
tests/local_testing/test_redis_increment_with_floor.py Normal file
View file

@ -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

264
tests/proxy_migration_tests/test_invalid_index_repair.py Normal file
View file

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

44
tests/test_litellm/caching/test_redis_cache.py
View file

@ -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)

187
tests/test_litellm/router_strategy/test_least_busy.py Normal file
View file

@ -0,0 +1,187 @@
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:
"""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.counts: dict[str, int] = {}
self.ttls: dict[str, int] = {}
def count(self, key: str) -> int | None:
return self.counts.get(key)
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:
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_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_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
return self.batch_get_counts(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_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"
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) == 1
assert shared.ttls == {key: 5}
@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 increment_with_floor(self, key: str, value: int, ttl: int) -> int:
raise ConnectionError("redis is down")
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
raise ConnectionError("redis is down")
@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.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
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
@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)
worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}})
worker.log_pre_api_call(model="m", messages=[], kwargs={})
assert shared.counts == {}

76
tests/test_litellm/router_strategy/test_router_routing_groups.py
View file

@ -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,81 @@ 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
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
# ---------------------------------------------------------------------------

77
tests/test_litellm/router_strategy/test_router_routing_plugins.py
View file

@ -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 [
{
@ -164,6 +176,71 @@ 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_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."""

8
type-discipline-budget.json
View file

@ -1,9 +1,9 @@
{
"LIT001": {
"limit": 22180
"limit": 22174
},
"LIT002": {
"limit": 26729
"limit": 26715
},
"LIT003": {
"limit": 261
@ -27,10 +27,10 @@
"limit": 0
},
"LIT010": {
"limit": 16426
"limit": 16398
},
"LIT011": {
"limit": 5506
"limit": 5504
},
"LIT012": {
"limit": 4486