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:
mateo-berri 2026-09-06 01:59:15 -07:00
parent 2f64272c9f
commit ef5f51abca
5 changed files with 121 additions and 1 deletions

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 \

View file

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

View file

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

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

View file

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