mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(least-busy): count for every router, and keep the expiry through a clamp
Two routers in one process shared a single handler, because the callback manager dedupes on the class name plus the handler's public attributes and the handler had none. The second router's requests were never counted. The handler now carries the id of the cache it was built on, so routers with different caches both register while the two selectors one router builds for its routing groups still collapse into one. Clamping a negative count back to zero used SET, which drops the key's TTL, so the next write started the hour over. It uses INCRBY by the negative amount now, which leaves the expiry alone. The Lua script had no test that ran it, so tests/local_testing covers both the sync and async paths against a real Redis, and the file is wired into the CircleCI job that provides one.
This commit is contained in:
parent
2f64272c9f
commit
ef5f51abca
5 changed files with 121 additions and 1 deletions
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -91,7 +91,7 @@ _BREAKER_GUARD_FRAME_NAMES: Final = frozenset(
|
|||
|
||||
_INCREMENT_WITH_FLOOR_LUA: Final = (
|
||||
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]) "
|
||||
"if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end "
|
||||
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end "
|
||||
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
|
||||
"return count"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -106,6 +106,7 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
|
||||
def __init__(self, router_cache: DualCache):
|
||||
self.router_cache = router_cache
|
||||
self.router_cache_id = str(id(router_cache))
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
|
||||
self._increment(kwargs, 1)
|
||||
|
|
|
|||
80
tests/local_testing/test_redis_increment_with_floor.py
Normal file
80
tests/local_testing/test_redis_increment_with_floor.py
Normal 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
|
||||
|
|
@ -473,6 +473,44 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch):
|
|||
assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0
|
||||
|
||||
|
||||
def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch):
|
||||
"""
|
||||
Least-busy hangs its counting off litellm's global callback lists, and those lists keep one
|
||||
logger per class unless the instances differ in a plain attribute. Two routers in one process
|
||||
(a second Router, or a per-request `user_config` one) therefore have to register separately:
|
||||
a second router whose selector is dropped counts nothing, reads zero for every deployment,
|
||||
and sends every request to whichever one is listed first.
|
||||
"""
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
|
||||
first = _build_router(routing_strategy="least-busy")
|
||||
second = _build_router(routing_strategy="least-busy")
|
||||
kwargs = {
|
||||
"litellm_params": {
|
||||
"metadata": {"model_group": "filtered-model"},
|
||||
"model_info": {"id": "deploy-1"},
|
||||
}
|
||||
}
|
||||
|
||||
for callback in litellm.input_callback:
|
||||
if isinstance(callback, CustomLogger):
|
||||
callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs)
|
||||
|
||||
assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1
|
||||
assert (
|
||||
second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"]
|
||||
== "deploy-2"
|
||||
)
|
||||
|
||||
for callback in litellm.callbacks:
|
||||
if isinstance(callback, CustomLogger):
|
||||
callback.log_success_event(kwargs, None, None, None)
|
||||
|
||||
assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0
|
||||
assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Direct helper coverage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue