mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
Merge 7e25fa409e into f32ac94ad1
This commit is contained in:
Krishnan
GitHub
commit
7209675163
6 changed files with 439 additions and 281 deletions
|
|
@ -56,6 +56,7 @@ if TYPE_CHECKING:
|
|||
from redis.asyncio import Redis, RedisCluster
|
||||
from redis.asyncio.client import Pipeline
|
||||
from redis.asyncio.cluster import ClusterPipeline
|
||||
from redis.client import Pipeline as SyncPipeline
|
||||
|
||||
from litellm._service_logger import ServiceLogging
|
||||
|
||||
|
|
@ -95,6 +96,16 @@ class _AsyncRedisCommands(Protocol):
|
|||
|
||||
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
|
||||
|
||||
def zrem(self, name: str, *values: str) -> Awaitable[int]: ...
|
||||
|
||||
|
||||
class _SyncRedisCommands(Protocol):
|
||||
"""Sync redis commands this cache issues; ``get_redis_client`` is untyped, so call sites need this protocol."""
|
||||
|
||||
def pipeline(self, transaction: bool = True) -> "SyncPipeline[bytes]": ...
|
||||
|
||||
def zrem(self, name: str, *values: str) -> int: ...
|
||||
|
||||
|
||||
_GENERIC_CALLER_MODULES: Final = frozenset(
|
||||
{
|
||||
|
|
@ -127,26 +138,12 @@ _GENERIC_CALLER_FRAME_NAMES: Final = frozenset(
|
|||
)
|
||||
_CALL_STACK_END_MODULES: Final = ("asyncio", "concurrent", "threading")
|
||||
|
||||
_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"
|
||||
)
|
||||
|
||||
_INCREMENT_IF_PRESENT_LUA: Final = (
|
||||
"if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"
|
||||
)
|
||||
|
||||
_LUA_COUNT: Final = TypeAdapter(int)
|
||||
_LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
|
||||
_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)
|
||||
)
|
||||
_COUNTS: Final = TypeAdapter(tuple[int, ...])
|
||||
|
||||
|
||||
def _is_generic_caller_frame(frame: FrameType) -> bool:
|
||||
|
|
@ -829,6 +826,9 @@ class RedisCache(BaseCache):
|
|||
def _async_commands(self) -> _AsyncRedisCommands:
|
||||
return self.init_async_client()
|
||||
|
||||
def _sync_commands(self) -> _SyncRedisCommands:
|
||||
return self.redis_client # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # get_redis_client is untyped
|
||||
|
||||
def check_and_fix_namespace(self, key: str) -> str:
|
||||
"""
|
||||
Make sure each key starts with the given namespace
|
||||
|
|
@ -954,41 +954,44 @@ 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.
|
||||
def track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None:
|
||||
"""Record ``member`` as in flight under ``key`` since ``started_at``, dropping members older than ``max_age``.
|
||||
|
||||
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.
|
||||
Every request is its own member of a sorted set, so finishing one request can only ever
|
||||
remove that request. A shared counter with a key-wide TTL forgot every request still in
|
||||
flight when it expired, and their later decrements made a busy deployment read idle. A
|
||||
member a dead worker never removed still stops counting ``max_age`` after it started.
|
||||
"""
|
||||
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)
|
||||
with self._sync_commands().pipeline(transaction=False) as pipe:
|
||||
pipe.zremrangebyscore(namespaced_key, "-inf", f"({started_at - max_age}")
|
||||
pipe.zadd(namespaced_key, {member: started_at})
|
||||
pipe.expire(namespaced_key, max_age)
|
||||
pipe.execute()
|
||||
|
||||
@_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.
|
||||
def untrack_in_flight(self, key: str, member: str) -> None:
|
||||
self._sync_commands().zrem(self.check_and_fix_namespace(key=key), member)
|
||||
|
||||
``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.
|
||||
@_redis_circuit_breaker_guard_sync
|
||||
def batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]:
|
||||
"""Count the members of each key that started at or after ``since``, in order.
|
||||
|
||||
Raises when Redis cannot answer, because a caller that falls back to its own numbers
|
||||
cannot tell a swallowed failure apart from "nothing is in flight".
|
||||
"""
|
||||
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))
|
||||
with self._sync_commands().pipeline(transaction=False) as pipe:
|
||||
for key in key_list:
|
||||
pipe.zcount(self.check_and_fix_namespace(key=key), since, "+inf")
|
||||
return _COUNTS.validate_python(pipe.execute())
|
||||
|
||||
@_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))
|
||||
async def async_batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]:
|
||||
"""Async twin of ``batch_count_in_flight``, raising on failure the same way."""
|
||||
async with self._async_commands().pipeline(transaction=False) as pipe:
|
||||
for key in key_list:
|
||||
pipe.zcount(self.check_and_fix_namespace(key=key), since, "+inf")
|
||||
return _COUNTS.validate_python(await pipe.execute())
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_scan_iter(self, pattern: str, count: int = 100) -> list:
|
||||
|
|
@ -1566,12 +1569,18 @@ class RedisCache(BaseCache):
|
|||
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()
|
||||
async def async_track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None:
|
||||
"""Async twin of ``track_in_flight``."""
|
||||
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 with self._async_commands().pipeline(transaction=False) as pipe:
|
||||
pipe.zremrangebyscore(namespaced_key, "-inf", f"({started_at - max_age}")
|
||||
pipe.zadd(namespaced_key, {member: started_at})
|
||||
pipe.expire(namespaced_key, max_age)
|
||||
await pipe.execute()
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_untrack_in_flight(self, key: str, member: str) -> None:
|
||||
await self._async_commands().zrem(self.check_and_fix_namespace(key=key), member)
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_increment_if_present(
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
import logging
|
||||
from collections.abc import Mapping, Sequence
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
|
@ -29,6 +32,7 @@ class _LitellmParams(TypedDict, total=False):
|
|||
|
||||
|
||||
class _CallKwargs(TypedDict, total=False):
|
||||
litellm_call_id: ReadOnly[str | None]
|
||||
litellm_params: ReadOnly[_LitellmParams | None]
|
||||
|
||||
|
||||
|
|
@ -45,11 +49,22 @@ _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
|
|||
_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CallRef:
|
||||
model_group: str
|
||||
deployment_id: str
|
||||
call_id: str | 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:
|
||||
def _in_flight_key(model_group: str, deployment_id: str) -> str:
|
||||
return f"{model_group}_in_flight:{deployment_id}"
|
||||
|
||||
|
||||
def _call_ref(kwargs: Mapping[str, object]) -> _CallRef | None:
|
||||
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||
return None
|
||||
try:
|
||||
|
|
@ -63,13 +78,12 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | 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)
|
||||
return _CallRef(model_group, str(deployment_id), call.get("litellm_call_id"))
|
||||
|
||||
|
||||
def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
|
||||
def _deployment_ids(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)
|
||||
str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -116,9 +130,24 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
logged_success: int = 0
|
||||
logged_failure: int = 0
|
||||
|
||||
def __init__(self, router_cache: DualCache):
|
||||
def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.time):
|
||||
self.router_cache = router_cache
|
||||
self.router_cache_id = str(id(router_cache))
|
||||
self._clock = clock
|
||||
self._worker_id = uuid.uuid4().hex
|
||||
|
||||
def _attempt_member(self, shared_key: str, call_id: str, delta: int) -> str:
|
||||
"""Name this call's sorted-set member, unique per concurrent attempt sharing a ``litellm_call_id``.
|
||||
|
||||
Callers can supply their own call id, so two concurrent calls may share one. Numbering the
|
||||
attempts per worker keeps one member per call, and finishing either call removes only one.
|
||||
"""
|
||||
attempts_key: Final = f"{shared_key}:{call_id}:attempts"
|
||||
memory: Final = self.router_cache.in_memory_cache
|
||||
after: Final = int(memory.increment_cache(attempts_key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)) # pyright: ignore[reportUnknownMemberType] # InMemoryCache.increment_cache leaves key untyped
|
||||
if after <= 0:
|
||||
memory.delete_cache(attempts_key)
|
||||
return f"{call_id}:{self._worker_id}:{after if delta > 0 else after + 1}"
|
||||
|
||||
@with_service_target("router_usage")
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
|
||||
|
|
@ -160,15 +189,19 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
def 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)
|
||||
deployment_ids: Final = _deployment_ids(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)))
|
||||
shared: Final = redis_cache.batch_count_in_flight(
|
||||
[_in_flight_key(model_group, deployment_id) for deployment_id in deployment_ids],
|
||||
self._clock() - IN_FLIGHT_COUNT_TTL_SECONDS,
|
||||
)
|
||||
except Exception as e:
|
||||
_warn_unreadable(model_group, e)
|
||||
else:
|
||||
return _least_busy(healthy_deployments, shared)
|
||||
keys: Final = tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids)
|
||||
local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
|
||||
return _least_busy(healthy_deployments, local)
|
||||
|
||||
|
|
@ -176,23 +209,27 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
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)
|
||||
deployment_ids: Final = _deployment_ids(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)))
|
||||
shared: Final = await redis_cache.async_batch_count_in_flight(
|
||||
[_in_flight_key(model_group, deployment_id) for deployment_id in deployment_ids],
|
||||
self._clock() - IN_FLIGHT_COUNT_TTL_SECONDS,
|
||||
)
|
||||
except Exception as e:
|
||||
_warn_unreadable(model_group, e)
|
||||
else:
|
||||
return _least_busy(healthy_deployments, shared)
|
||||
keys: Final = tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids)
|
||||
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)
|
||||
ref: Final = _call_ref(kwargs)
|
||||
if ref is None:
|
||||
return
|
||||
key: Final = _request_count_key(*ref)
|
||||
key: Final = _request_count_key(ref.model_group, ref.deployment_id)
|
||||
redis_cache: Final = self.router_cache.redis_cache
|
||||
try:
|
||||
local: Final = self.router_cache.increment_cache(
|
||||
|
|
@ -200,17 +237,22 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
)
|
||||
if local < 0:
|
||||
self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
if redis_cache is None:
|
||||
if redis_cache is None or ref.call_id is None:
|
||||
return
|
||||
redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
shared_key: Final = _in_flight_key(ref.model_group, ref.deployment_id)
|
||||
member: Final = self._attempt_member(shared_key, ref.call_id, delta)
|
||||
if delta > 0:
|
||||
redis_cache.track_in_flight(shared_key, member, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
else:
|
||||
redis_cache.untrack_in_flight(shared_key, member)
|
||||
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)
|
||||
ref: Final = _call_ref(kwargs)
|
||||
if ref is None:
|
||||
return
|
||||
key: Final = _request_count_key(*ref)
|
||||
key: Final = _request_count_key(ref.model_group, ref.deployment_id)
|
||||
redis_cache: Final = self.router_cache.redis_cache
|
||||
try:
|
||||
local: Final = await self.router_cache.async_increment_cache(
|
||||
|
|
@ -218,8 +260,13 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
)
|
||||
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:
|
||||
if redis_cache is None or ref.call_id is None:
|
||||
return
|
||||
await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
shared_key: Final = _in_flight_key(ref.model_group, ref.deployment_id)
|
||||
member: Final = self._attempt_member(shared_key, ref.call_id, delta)
|
||||
if delta > 0:
|
||||
await redis_cache.async_track_in_flight(shared_key, member, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
else:
|
||||
await redis_cache.async_untrack_in_flight(shared_key, member)
|
||||
except Exception as e:
|
||||
_warn_unwritable(key, e)
|
||||
|
|
|
|||
61
tests/integration/routing/test_redis_in_flight_tracking.py
Normal file
61
tests/integration/routing/test_redis_in_flight_tracking.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
import os
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
MAX_AGE: Final = 600
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def in_flight() -> Iterator[tuple[RedisCache, str, str]]:
|
||||
cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
key: Final = f"in-flight-{uuid.uuid4()}"
|
||||
yield cache, key, cache.check_and_fix_namespace(key=key)
|
||||
cache.delete_cache(key)
|
||||
|
||||
|
||||
def test_requests_finishing_after_newer_ones_started_leave_the_newer_ones_counted(
|
||||
in_flight: tuple[RedisCache, str, str],
|
||||
) -> None:
|
||||
cache, key, _ = in_flight
|
||||
now: Final = time.time()
|
||||
|
||||
for call_id in ("old-1", "old-2", "old-3", "new-1", "new-2", "new-3"):
|
||||
cache.track_in_flight(key, call_id, now, MAX_AGE)
|
||||
for call_id in ("old-1", "old-2", "old-3"):
|
||||
cache.untrack_in_flight(key, call_id)
|
||||
cache.track_in_flight(key, "new-4", now, MAX_AGE)
|
||||
|
||||
assert cache.batch_count_in_flight([key], since=now - MAX_AGE) == (4,)
|
||||
|
||||
|
||||
def test_a_request_older_than_max_age_ages_out_without_touching_the_rest(
|
||||
in_flight: tuple[RedisCache, str, str],
|
||||
) -> None:
|
||||
cache, key, namespaced_key = in_flight
|
||||
now: Final = time.time()
|
||||
|
||||
cache.track_in_flight(key, "orphan", now - MAX_AGE - 1, MAX_AGE)
|
||||
cache.track_in_flight(key, "running", now - MAX_AGE + 60, MAX_AGE)
|
||||
cache.track_in_flight(key, "new", now, MAX_AGE)
|
||||
|
||||
assert cache.batch_count_in_flight([key], since=now - MAX_AGE) == (2,)
|
||||
assert cache.redis_client.zcard(namespaced_key) == 2
|
||||
assert 0 < cache.redis_client.ttl(namespaced_key) <= MAX_AGE
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_async_twins_share_the_same_state(in_flight: tuple[RedisCache, str, str]) -> None:
|
||||
cache, key, _ = in_flight
|
||||
now: Final = time.time()
|
||||
|
||||
await cache.async_track_in_flight(key, "async-1", now, MAX_AGE)
|
||||
cache.track_in_flight(key, "sync-1", now, MAX_AGE)
|
||||
await cache.async_untrack_in_flight(key, "sync-1")
|
||||
|
||||
assert await cache.async_batch_count_in_flight([key], since=now - MAX_AGE) == (1,)
|
||||
|
|
@ -1,69 +0,0 @@
|
|||
import os
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
TTL: Final = 600
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def counter() -> Iterator[tuple[RedisCache, str, str]]:
|
||||
cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
key: Final = f"increment-with-floor-{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: tuple[RedisCache, str, str]) -> None:
|
||||
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: tuple[RedisCache, str, str]) -> None:
|
||||
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: tuple[RedisCache, str, str]) -> None:
|
||||
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: tuple[RedisCache, str, str]) -> None:
|
||||
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: tuple[RedisCache, str, str]) -> None:
|
||||
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
|
||||
|
|
@ -506,48 +506,92 @@ def test_sync_get_cache_failure_feeds_the_breaker_and_logs_a_well_formed_record(
|
|||
sync_batch_redis_cache.get_cache("lit7468")
|
||||
|
||||
|
||||
def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache):
|
||||
def test_batch_count_in_flight_raises_where_batch_get_cache_reports_a_miss(redis_no_ping: None):
|
||||
"""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
|
||||
apart from "nothing is in flight". 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"]) == {}
|
||||
from fakeredis import FakeRedis, FakeServer
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
with pytest.raises(OSError, match="redis unavailable"):
|
||||
sync_batch_redis_cache.batch_get_counts(["lit7039"])
|
||||
unreachable = FakeServer()
|
||||
unreachable.connected = False
|
||||
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
"litellm._redis.get_redis_client", return_value=FakeRedis(server=unreachable)
|
||||
):
|
||||
cache = RedisCache(host="127.0.0.1", port=6379)
|
||||
|
||||
assert cache.batch_get_cache(key_list=["lit7039"]) == {}
|
||||
|
||||
with pytest.raises(RedisConnectionError):
|
||||
cache.batch_count_in_flight(["lit7039"], since=0.0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def in_flight_redis() -> Iterator[RedisCache]:
|
||||
from fakeredis import FakeRedis, FakeServer
|
||||
from fakeredis.aioredis import FakeRedis as AsyncFakeRedis
|
||||
|
||||
server = FakeServer()
|
||||
with (
|
||||
patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")),
|
||||
patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
"litellm._redis.get_redis_client", return_value=FakeRedis(server=server)
|
||||
),
|
||||
):
|
||||
cache = RedisCache(host="127.0.0.1", port=6379, namespace="ns")
|
||||
with patch.object(cache, "init_async_client", return_value=AsyncFakeRedis(server=server)):
|
||||
yield cache
|
||||
|
||||
|
||||
IN_FLIGHT_START: Final = 1_700_000_000.0
|
||||
IN_FLIGHT_MAX_AGE: Final = 600
|
||||
|
||||
|
||||
def test_finishing_a_request_removes_only_that_request(in_flight_redis: RedisCache):
|
||||
for call_id in ("old-1", "old-2", "old-3", "new-1"):
|
||||
in_flight_redis.track_in_flight("dep-a", call_id, IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
for call_id in ("old-1", "old-2", "old-3", "old-3", "never-started"):
|
||||
in_flight_redis.untrack_in_flight("dep-a", call_id)
|
||||
|
||||
assert in_flight_redis.batch_count_in_flight(["dep-a"], since=IN_FLIGHT_START) == (1,)
|
||||
|
||||
|
||||
def test_in_flight_counts_come_back_in_key_order_with_unknown_keys_as_zero(in_flight_redis: RedisCache):
|
||||
in_flight_redis.track_in_flight("dep-a", "a-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
in_flight_redis.track_in_flight("dep-c", "c-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
in_flight_redis.track_in_flight("dep-c", "c-2", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
|
||||
assert in_flight_redis.batch_count_in_flight(["dep-c", "dep-b", "dep-a"], since=IN_FLIGHT_START) == (2, 0, 1)
|
||||
|
||||
|
||||
def test_requests_that_started_before_the_cutoff_are_not_counted(in_flight_redis: RedisCache):
|
||||
in_flight_redis.track_in_flight("dep-a", "older", IN_FLIGHT_START - 1, IN_FLIGHT_MAX_AGE)
|
||||
in_flight_redis.track_in_flight("dep-a", "at-cutoff", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
|
||||
assert in_flight_redis.batch_count_in_flight(["dep-a"], since=IN_FLIGHT_START) == (1,)
|
||||
|
||||
|
||||
def test_tracking_a_request_drops_members_older_than_max_age(in_flight_redis: RedisCache):
|
||||
in_flight_redis.track_in_flight("dep-a", "orphan", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
in_flight_redis.track_in_flight("dep-a", "edge", IN_FLIGHT_START + 1, IN_FLIGHT_MAX_AGE)
|
||||
in_flight_redis.track_in_flight("dep-a", "new", IN_FLIGHT_START + 1 + IN_FLIGHT_MAX_AGE, IN_FLIGHT_MAX_AGE)
|
||||
|
||||
assert in_flight_redis.redis_client.zrange("ns:dep-a", 0, -1) == [b"edge", b"new"]
|
||||
assert 0 < in_flight_redis.redis_client.ttl("ns:dep-a") <= IN_FLIGHT_MAX_AGE
|
||||
|
||||
|
||||
@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)
|
||||
async def test_async_in_flight_tracking_shares_the_sync_state(in_flight_redis: RedisCache):
|
||||
await in_flight_redis.async_track_in_flight("dep-a", "async-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
in_flight_redis.track_in_flight("dep-a", "sync-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
await in_flight_redis.async_track_in_flight("dep-b", "async-2", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE)
|
||||
await in_flight_redis.async_untrack_in_flight("dep-a", "sync-1")
|
||||
|
||||
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)
|
||||
assert await in_flight_redis.async_batch_count_in_flight(["dep-a", "dep-b"], since=IN_FLIGHT_START) == (1, 1)
|
||||
assert in_flight_redis.batch_count_in_flight(["dep-a", "dep-b"], since=IN_FLIGHT_START) == (1, 1)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -1033,8 +1077,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke
|
|||
monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"])
|
||||
gc.collect()
|
||||
ServiceLogging()
|
||||
assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == (
|
||||
redis_breaker_state_counts()["closed"]
|
||||
assert (
|
||||
REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"})
|
||||
== (redis_breaker_state_counts()["closed"])
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,19 +1,24 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import fakeredis
|
||||
import fakeredis.aioredis
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
|
||||
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
|
@ -21,102 +26,172 @@ 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]
|
||||
START: Final = 1_700_000_000.0
|
||||
|
||||
|
||||
def _call_kwargs(deployment_id: str) -> dict[str, object]:
|
||||
return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}}
|
||||
def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]:
|
||||
return {
|
||||
"litellm_call_id": call_id,
|
||||
"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."""
|
||||
|
||||
class Clock:
|
||||
def __init__(self) -> None:
|
||||
self.counts: dict[str, int] = {}
|
||||
self.ttls: dict[str, int] = {}
|
||||
self.now = START
|
||||
|
||||
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 __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
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.fixture
|
||||
def redis_server() -> fakeredis.FakeServer:
|
||||
return fakeredis.FakeServer()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def shared_redis(redis_server: fakeredis.FakeServer) -> Iterator[RedisCache]:
|
||||
server: Final = redis_server
|
||||
with (
|
||||
patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")),
|
||||
patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
"litellm._redis.get_redis_client", return_value=fakeredis.FakeRedis(server=server)
|
||||
),
|
||||
):
|
||||
cache: Final = RedisCache(host="127.0.0.1", port=6379)
|
||||
with patch.object(cache, "init_async_client", return_value=fakeredis.aioredis.FakeRedis(server=server)):
|
||||
yield cache
|
||||
|
||||
|
||||
class UnreachableRedis:
|
||||
def __init__(self, error: Exception) -> None:
|
||||
self.error = error
|
||||
|
||||
def track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None:
|
||||
raise self.error
|
||||
|
||||
async def async_track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None:
|
||||
raise self.error
|
||||
|
||||
def untrack_in_flight(self, key: str, member: str) -> None:
|
||||
raise self.error
|
||||
|
||||
async def async_untrack_in_flight(self, key: str, member: str) -> None:
|
||||
raise self.error
|
||||
|
||||
def batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]:
|
||||
raise self.error
|
||||
|
||||
async def async_batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]:
|
||||
raise self.error
|
||||
|
||||
|
||||
def _worker(
|
||||
shared: RedisCache | UnreachableRedis | None, clock: Callable[[], float] = time.time
|
||||
) -> LeastBusyLoggingHandler:
|
||||
cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # UnreachableRedis is a duck-typed Redis double
|
||||
return LeastBusyLoggingHandler(router_cache=cache, clock=clock)
|
||||
|
||||
|
||||
@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)
|
||||
async def test_worker_routes_around_a_request_another_worker_started(shared_redis: RedisCache) -> None:
|
||||
streaming_worker: Final = _worker(shared_redis)
|
||||
picking_worker: Final = _worker(shared_redis)
|
||||
|
||||
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)
|
||||
picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "quick"))
|
||||
await picking_worker.async_log_success_event(_call_kwargs("dep-a", "quick"), None, None, None)
|
||||
|
||||
streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "stream"))
|
||||
|
||||
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)
|
||||
await streaming_worker.async_log_success_event(_call_kwargs("dep-a", "stream"), 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)
|
||||
def test_sync_pick_reads_the_shared_counts(shared_redis: RedisCache) -> None:
|
||||
streaming_worker: Final = _worker(shared_redis)
|
||||
picking_worker: Final = _worker(shared_redis)
|
||||
|
||||
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)
|
||||
picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "quick"))
|
||||
picking_worker.log_success_event(_call_kwargs("dep-a", "quick"), None, None, None)
|
||||
|
||||
streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "stream"))
|
||||
|
||||
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None)
|
||||
streaming_worker.log_failure_event(_call_kwargs("dep-a", "stream"), 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"
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_request_outliving_the_ttl_cannot_cancel_a_newer_request_when_it_finishes(
|
||||
shared_redis: RedisCache,
|
||||
) -> None:
|
||||
"""A shared counter forgot every in-flight request when its key expired, then each late
|
||||
decrement cancelled a request that was really running, so a busy deployment read idle and
|
||||
kept getting a full share of new traffic (#48994)"""
|
||||
clock: Final = Clock()
|
||||
worker: Final = _worker(shared_redis, clock)
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "old"))
|
||||
clock.now = START + IN_FLIGHT_COUNT_TTL_SECONDS + 1
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "new"))
|
||||
await worker.async_log_success_event(_call_kwargs("dep-a", "old"), None, None, None)
|
||||
|
||||
assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS}
|
||||
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
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_finishing_one_of_several_calls_sharing_a_call_id_keeps_the_others_counted(
|
||||
shared_redis: RedisCache,
|
||||
) -> None:
|
||||
"""Callers can pass their own litellm_call_id, so concurrent calls may share one. Tracking the
|
||||
id itself merged them into one entry, and the first to finish hid the rest from other workers"""
|
||||
worker: Final = _worker(shared_redis)
|
||||
other_worker: Final = _worker(shared_redis)
|
||||
picking_worker: Final = _worker(shared_redis)
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared"))
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared"))
|
||||
other_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared"))
|
||||
picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "b"))
|
||||
await worker.async_log_success_event(_call_kwargs("dep-a", "shared"), None, None, None)
|
||||
|
||||
assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
worker.log_success_event(_call_kwargs("dep-a", "shared"), None, None, None)
|
||||
await other_worker.async_log_failure_event(_call_kwargs("dep-a", "shared"), None, None, None)
|
||||
|
||||
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_request_nobody_finished_stops_counting_once_it_is_older_than_the_ttl(
|
||||
shared_redis: RedisCache,
|
||||
) -> None:
|
||||
"""A worker that dies mid-request never reports the end. Its request has to age out on its
|
||||
own, without taking the deployment's other in-flight requests with it"""
|
||||
clock: Final = Clock()
|
||||
worker: Final = _worker(shared_redis, clock)
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "orphan"))
|
||||
clock.now = START + IN_FLIGHT_COUNT_TTL_SECONDS - 1
|
||||
|
||||
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "running-b"))
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "running-a"))
|
||||
clock.now = START + IN_FLIGHT_COUNT_TTL_SECONDS + 1
|
||||
|
||||
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
await worker.async_log_success_event(_call_kwargs("dep-b", "running-b"), None, None, None)
|
||||
|
||||
assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -134,17 +209,9 @@ async def test_counts_stay_in_memory_without_redis() -> None:
|
|||
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: Final = _worker(UnreachableRedis(ConnectionError("redis is down")))
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
|
||||
|
|
@ -156,22 +223,6 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> 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)
|
||||
|
|
@ -188,27 +239,20 @@ async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() ->
|
|||
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)
|
||||
def test_calls_without_a_deployment_are_ignored(shared_redis: RedisCache, redis_server: fakeredis.FakeServer) -> None:
|
||||
worker: Final = _worker(shared_redis)
|
||||
|
||||
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 == {}
|
||||
|
||||
|
||||
class OpenBreakerRedis(SharedRedisCounters):
|
||||
def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
|
||||
raise RedisCircuitBreakerOpenError("Redis circuit breaker is open")
|
||||
|
||||
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
|
||||
raise RedisCircuitBreakerOpenError("Redis circuit breaker is open")
|
||||
assert fakeredis.FakeRedis(server=redis_server).dbsize() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None:
|
||||
worker: Final = _worker(OpenBreakerRedis())
|
||||
async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
worker: Final = _worker(UnreachableRedis(RedisCircuitBreakerOpenError("Redis circuit breaker is open")))
|
||||
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"):
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
|
|
@ -307,6 +351,12 @@ def _chat_stream(_: httpx.Request) -> httpx.Response:
|
|||
)
|
||||
|
||||
|
||||
async def _open_chat_stream(router: Router) -> CustomStreamWrapper:
|
||||
return await router.acompletion(
|
||||
model=ROUTER_GROUP, messages=[{"role": "user", "content": "write a poem"}], stream=True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_spreads_open_chat_streams_across_least_busy_deployments(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
|
|
@ -315,12 +365,7 @@ async def test_router_spreads_open_chat_streams_across_least_busy_deployments(
|
|||
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_chat_stream)
|
||||
router: Final = _least_busy_router()
|
||||
|
||||
streams: Final = [
|
||||
await router.acompletion(
|
||||
model=ROUTER_GROUP, messages=[{"role": "user", "content": "write a poem"}], stream=True
|
||||
)
|
||||
for _ in ROUTER_DEPLOYMENT_IDS
|
||||
]
|
||||
streams: Final = [await _open_chat_stream(router) for _ in ROUTER_DEPLOYMENT_IDS]
|
||||
|
||||
assert _router_counts(router) == {"1": 1, "2": 1, "3": 1}
|
||||
assert sorted(stream._hidden_params["model_id"] for stream in streams) == list(ROUTER_DEPLOYMENT_IDS)
|
||||
|
|
@ -379,9 +424,7 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts
|
|||
if use_async
|
||||
else router.get_available_deployment(model=ROUTER_GROUP, messages=None)
|
||||
)
|
||||
response: Final = await router.acompletion(
|
||||
model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}])
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
await GLOBAL_LOGGING_WORKER.flush()
|
||||
|
|
@ -391,3 +434,25 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts
|
|||
assert response._hidden_params["model_id"] == "1"
|
||||
assert route.calls.last.request.headers["authorization"] == "Bearer key-1"
|
||||
assert _router_counts(router) == seeded
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_spreads_open_streams_across_workers_through_the_shared_counts(
|
||||
shared_redis: RedisCache, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
keys_used: Final[list[str]] = []
|
||||
|
||||
def record_key_and_stream(request: httpx.Request) -> httpx.Response:
|
||||
keys_used.append(request.headers["authorization"])
|
||||
return _chat_stream(request)
|
||||
|
||||
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=record_key_and_stream)
|
||||
routers: Final = tuple(_least_busy_router() for _ in ROUTER_DEPLOYMENT_IDS)
|
||||
for router in routers:
|
||||
router.cache.redis_cache = shared_redis
|
||||
|
||||
for router in routers:
|
||||
await _open_chat_stream(router)
|
||||
|
||||
assert sorted(keys_used) == [f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue