mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(router): count concurrent successes atomically in usage-based routing
Each success read the model group's {deployment: count} dict, added to it and wrote it back. With Redis, concurrent callbacks interleave between the read and the write and overwrite each other, so TPM/RPM were undercounted. Each deployment now has its own counter updated with a single atomic increment; the in-memory dict routing reads is filled from those totals. Fixes #44917
This commit is contained in:
Charan Kotra
parent
dddbaa803d
commit
2e5c5611c0
2 changed files with 134 additions and 40 deletions
|
|
@ -1,6 +1,7 @@
|
|||
#### What this does ####
|
||||
# identifies lowest tpm deployment
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -28,6 +29,40 @@ class LowestTPMLoggingHandler(CustomLogger):
|
|||
self.router_cache = router_cache
|
||||
self.routing_args = RoutingArgs(**routing_args)
|
||||
|
||||
def _cached_usage(self, usage_key: str) -> Mapping[str, float]:
|
||||
"""The {deployment_id: count} dict this process has recorded for a model group this minute."""
|
||||
cached: Final = self.router_cache.get_cache(key=usage_key, local_only=True)
|
||||
if not isinstance(cached, Mapping):
|
||||
return {}
|
||||
return {k: float(v) for k, v in cached.items() if isinstance(k, str) and isinstance(v, (int, float))}
|
||||
|
||||
def _record_usage(self, usage_key: str, deployment_id: str, total: float | None) -> None:
|
||||
"""Copy a deployment's running total into the in-memory {deployment_id: count} dict for its
|
||||
model group, which routing and the router's pre-call RPM check read. Nothing awaits between
|
||||
the read and the write, so concurrent callbacks in this process can't overwrite each other."""
|
||||
if total is None:
|
||||
return
|
||||
recorded: Final = self._cached_usage(usage_key)
|
||||
usage: Final = {**recorded, deployment_id: max(recorded.get(deployment_id, 0.0), total)}
|
||||
self.router_cache.set_cache(key=usage_key, value=usage, local_only=True, ttl=self.routing_args.ttl)
|
||||
|
||||
def _get_usage(self, usage_key: str, deployment_ids: Sequence[str]) -> Mapping[str, float]:
|
||||
"""This process's recorded usage, plus the shared counter for any deployment it hasn't served
|
||||
this minute (read from Redis, if there is one, so other instances' traffic counts).
|
||||
|
||||
Counters are kept per deployment, f"{usage_key}:{deployment_id}", so each success is one
|
||||
atomic increment (INCRBYFLOAT in Redis) instead of a read-modify-write of a shared dict,
|
||||
which loses updates when callbacks run concurrently."""
|
||||
recorded: Final = self._cached_usage(usage_key)
|
||||
unseen: Final = tuple(i for i in deployment_ids if i not in recorded)
|
||||
if not unseen:
|
||||
return recorded
|
||||
counts: Final = self.router_cache.batch_get_cache(keys=[f"{usage_key}:{i}" for i in unseen])
|
||||
if not isinstance(counts, Sequence):
|
||||
return recorded
|
||||
shared: Final = {i: float(c) for i, c in zip(unseen, counts) if c is not None}
|
||||
return {**shared, **recorded}
|
||||
|
||||
@with_service_target("router_usage")
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||
|
|
@ -62,17 +97,20 @@ class LowestTPMLoggingHandler(CustomLogger):
|
|||
# Update usage
|
||||
# ------------
|
||||
|
||||
## TPM
|
||||
request_count_dict = self.router_cache.get_cache(key=tpm_key) or {}
|
||||
request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens
|
||||
if not isinstance(id, str) or not isinstance(total_tokens, int):
|
||||
return
|
||||
|
||||
self.router_cache.set_cache(key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl)
|
||||
## TPM
|
||||
tpm_total: Final = self.router_cache.increment_cache(
|
||||
key=f"{tpm_key}:{id}", value=total_tokens, ttl=self.routing_args.ttl
|
||||
)
|
||||
self._record_usage(tpm_key, id, tpm_total)
|
||||
|
||||
## RPM
|
||||
request_count_dict = self.router_cache.get_cache(key=rpm_key) or {}
|
||||
request_count_dict[id] = request_count_dict.get(id, 0) + 1
|
||||
|
||||
self.router_cache.set_cache(key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl)
|
||||
rpm_total: Final = self.router_cache.increment_cache(
|
||||
key=f"{rpm_key}:{id}", value=1, ttl=self.routing_args.ttl
|
||||
)
|
||||
self._record_usage(rpm_key, id, rpm_total)
|
||||
|
||||
### TESTING ###
|
||||
if self.test_flag:
|
||||
|
|
@ -123,21 +161,20 @@ class LowestTPMLoggingHandler(CustomLogger):
|
|||
# ------------
|
||||
# update cache
|
||||
|
||||
## TPM
|
||||
request_count_dict = await self.router_cache.async_get_cache(key=tpm_key) or {}
|
||||
request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens
|
||||
if not isinstance(id, str) or not isinstance(total_tokens, (int, float)):
|
||||
return
|
||||
|
||||
await self.router_cache.async_set_cache(
|
||||
key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl
|
||||
## TPM
|
||||
tpm_total: Final = await self.router_cache.async_increment_cache(
|
||||
key=f"{tpm_key}:{id}", value=total_tokens, ttl=self.routing_args.ttl
|
||||
)
|
||||
self._record_usage(tpm_key, id, tpm_total)
|
||||
|
||||
## RPM
|
||||
request_count_dict = await self.router_cache.async_get_cache(key=rpm_key) or {}
|
||||
request_count_dict[id] = request_count_dict.get(id, 0) + 1
|
||||
|
||||
await self.router_cache.async_set_cache(
|
||||
key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl
|
||||
rpm_total: Final = await self.router_cache.async_increment_cache(
|
||||
key=f"{rpm_key}:{id}", value=1, ttl=self.routing_args.ttl
|
||||
)
|
||||
self._record_usage(rpm_key, id, rpm_total)
|
||||
|
||||
### TESTING ###
|
||||
if self.test_flag:
|
||||
|
|
@ -169,8 +206,9 @@ class LowestTPMLoggingHandler(CustomLogger):
|
|||
tpm_key: Final = f"{model_group}:tpm:{current_minute}"
|
||||
rpm_key: Final = f"{model_group}:rpm:{current_minute}"
|
||||
|
||||
tpm_dict = self.router_cache.get_cache(key=tpm_key)
|
||||
rpm_dict: Final = self.router_cache.get_cache(key=rpm_key)
|
||||
deployment_ids: Final = tuple(str(d["model_info"]["id"]) for d in healthy_deployments)
|
||||
tpm_dict: Final = dict(self._get_usage(tpm_key, deployment_ids))
|
||||
rpm_dict: Final = self._get_usage(rpm_key, deployment_ids)
|
||||
|
||||
verbose_router_logger.debug("tpm_key=%s, tpm_dict: %s, rpm_dict: %s", tpm_key, tpm_dict, rpm_dict)
|
||||
try:
|
||||
|
|
@ -183,16 +221,8 @@ class LowestTPMLoggingHandler(CustomLogger):
|
|||
# ----------------------
|
||||
lowest_tpm = float("inf")
|
||||
|
||||
if tpm_dict is None: # base case - none of the deployments have been used
|
||||
# initialize a tpm dict with {model_id: 0}
|
||||
tpm_dict = {}
|
||||
for deployment in healthy_deployments:
|
||||
tpm_dict[deployment["model_info"]["id"]] = 0
|
||||
else:
|
||||
for d in healthy_deployments:
|
||||
## if healthy deployment not yet used
|
||||
if d["model_info"]["id"] not in tpm_dict:
|
||||
tpm_dict[d["model_info"]["id"]] = 0
|
||||
for deployment_id in deployment_ids: ## if healthy deployment not yet used
|
||||
_ = tpm_dict.setdefault(deployment_id, 0)
|
||||
|
||||
all_deployments: Final = tpm_dict
|
||||
|
||||
|
|
@ -229,7 +259,7 @@ class LowestTPMLoggingHandler(CustomLogger):
|
|||
|
||||
if (
|
||||
item_tpm + input_tokens > _deployment_tpm
|
||||
or (rpm_dict is not None and item in rpm_dict)
|
||||
or (item in rpm_dict)
|
||||
and (rpm_dict[item] + 1 >= _deployment_rpm)
|
||||
):
|
||||
continue
|
||||
|
|
|
|||
|
|
@ -1,12 +1,16 @@
|
|||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
|
||||
from litellm import Router, token_counter
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler
|
||||
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage
|
||||
from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict, RouterErrors
|
||||
|
||||
|
|
@ -43,14 +47,11 @@ def test_usage_based_routing_v1_selects_the_lowest_recorded_tpm() -> None:
|
|||
}
|
||||
now: Final = datetime.now()
|
||||
cache_keys: Final = tuple(
|
||||
f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}"
|
||||
for offset in range(60)
|
||||
f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" for offset in range(60)
|
||||
)
|
||||
|
||||
for cache_key in cache_keys:
|
||||
router.cache.set_cache(
|
||||
key=cache_key, value=usage_by_deployment, ttl=float("inf")
|
||||
)
|
||||
router.cache.set_cache(key=cache_key, value=usage_by_deployment, ttl=float("inf"))
|
||||
|
||||
deployment: Final = router.get_available_deployment(
|
||||
model=MODEL_GROUP,
|
||||
|
|
@ -135,9 +136,7 @@ async def test_v2_subclass_overriding_async_get_available_deployments_with_the_o
|
|||
)
|
||||
router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache, routing_args={})
|
||||
|
||||
response: Final = await router.acompletion(
|
||||
model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}]
|
||||
)
|
||||
response: Final = await router.acompletion(model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}])
|
||||
|
||||
assert response.choices[0].message.content in {
|
||||
f"from {HIGH_USAGE_DEPLOYMENT_ID}",
|
||||
|
|
@ -209,3 +208,68 @@ def test_usage_based_routing_v1_rejects_a_deployment_that_reached_its_recorded_r
|
|||
router: Final = _router_with_recorded_rpm(recorded=RPM_LIMIT, enable_pre_call_checks=enable_pre_call_checks)
|
||||
with pytest.raises(ValueError, match=RouterErrors.no_deployments_available.value):
|
||||
router.completion(model=MODEL_GROUP, messages=[{"role": "user", "content": "hello"}])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- concurrent successes, Redis (#44917)
|
||||
class _RedisWithReplyLatency(fakeredis.FakeAsyncRedis):
|
||||
"""fakeredis replies instantly, so concurrent callers never overlap. A real Redis reads the value
|
||||
when the command arrives and the reply then spends time on the network while other callers run.
|
||||
Yielding after the read reproduces that window."""
|
||||
|
||||
async def get(self, *args, **kwargs):
|
||||
value = await super().get(*args, **kwargs)
|
||||
await asyncio.sleep(0)
|
||||
return value
|
||||
|
||||
|
||||
def _limited_deployment(deployment_id: str, **limits: int) -> DeploymentTypedDict:
|
||||
deployment: Final = _deployment(deployment_id)
|
||||
return {**deployment, "litellm_params": {**deployment["litellm_params"], **limits}}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_backed_handler():
|
||||
with patch("litellm._redis.get_redis_client", return_value=MagicMock()):
|
||||
redis_cache: Final = RedisCache(host="localhost", port=6379)
|
||||
with patch.object(redis_cache, "init_async_client", return_value=_RedisWithReplyLatency()):
|
||||
yield LowestTPMLoggingHandler(router_cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
|
||||
|
||||
|
||||
async def _record_successes(
|
||||
handler: LowestTPMLoggingHandler, deployment_id: str, total_tokens: int, times: int = 1
|
||||
) -> None:
|
||||
kwargs: Final = {"litellm_params": {"metadata": {"model_group": MODEL_GROUP}, "model_info": {"id": deployment_id}}}
|
||||
response: Final = {"usage": {"total_tokens": total_tokens}}
|
||||
await asyncio.gather(*(handler.async_log_success_event(kwargs, response, None, None) for _ in range(times)))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_based_routing_v1_counts_every_concurrent_success_toward_tpm(redis_backed_handler) -> None:
|
||||
await _record_successes(redis_backed_handler, LOW_USAGE_DEPLOYMENT_ID, total_tokens=3, times=20)
|
||||
await _record_successes(redis_backed_handler, HIGH_USAGE_DEPLOYMENT_ID, total_tokens=40)
|
||||
|
||||
deployment: Final = redis_backed_handler.get_available_deployments(
|
||||
model_group=MODEL_GROUP,
|
||||
healthy_deployments=[
|
||||
_limited_deployment(LOW_USAGE_DEPLOYMENT_ID, tpm=50),
|
||||
_limited_deployment(HIGH_USAGE_DEPLOYMENT_ID, tpm=1000),
|
||||
],
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] == HIGH_USAGE_DEPLOYMENT_ID
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_based_routing_v1_counts_every_concurrent_success_toward_rpm(redis_backed_handler) -> None:
|
||||
await _record_successes(redis_backed_handler, LOW_USAGE_DEPLOYMENT_ID, total_tokens=1, times=20)
|
||||
await _record_successes(redis_backed_handler, HIGH_USAGE_DEPLOYMENT_ID, total_tokens=40)
|
||||
|
||||
deployment: Final = redis_backed_handler.get_available_deployments(
|
||||
model_group=MODEL_GROUP,
|
||||
healthy_deployments=[
|
||||
_limited_deployment(LOW_USAGE_DEPLOYMENT_ID, rpm=20),
|
||||
_limited_deployment(HIGH_USAGE_DEPLOYMENT_ID, rpm=1000),
|
||||
],
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] == HIGH_USAGE_DEPLOYMENT_ID
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue