mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
b1b236826a
* feat(redis): route the router, rate limiters and auth through the native store families * style(redis): satisfy the python lint gates Format the rate limit store, sort its __all__ and give the routing read batch and two tests concrete list and dict types. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(router): give the registry's default usage store a memory-only handle RouterUsageStore now takes a handle, so the default passed None and any usage read before RustCore swapped in its store raised AttributeError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(redis): build stores from native handles and wait for router tasks Tests still used the pre-refactor store constructors (Store(None), handle= factories, NativeIncrement, tuple scopes). The router outcome tests counted 20 loop ticks, too few now that the allowed-fails increment resolves on the native runtime, so settle() waits for the scheduled tasks instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
102 lines
4.1 KiB
Python
102 lines
4.1 KiB
Python
import asyncio
|
|
from collections.abc import Iterator, Sequence
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from typing import TYPE_CHECKING, Final
|
|
|
|
from litellm._logging import verbose_router_logger
|
|
from litellm.router_deprecated import renamed_router_param
|
|
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage
|
|
from litellm.router_utils import routing_prefetch
|
|
from litellm.router_utils.routing_prefetch import TierRead, selector_read
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span
|
|
|
|
from litellm.router_api import RoutingContext
|
|
|
|
__all__ = ("RoutingPrefetch", "RoutingReadBatch")
|
|
|
|
|
|
def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: Sequence[object]) -> list[str]:
|
|
tpm_keys, rpm_keys = selector.usage_counter_keys(list(deployments))
|
|
return tpm_keys + rpm_keys
|
|
|
|
|
|
def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: Sequence[object]) -> TierRead | None:
|
|
return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments))
|
|
|
|
|
|
class RoutingPrefetch:
|
|
@staticmethod
|
|
@renamed_router_param("litellm_router_instance")
|
|
def arm(
|
|
ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: Sequence[object]
|
|
) -> None:
|
|
RoutingPrefetch.arm_cooldowns(ctx)
|
|
read: Final = _selector_read(usage_selector, deployments)
|
|
if read is not None:
|
|
routing_prefetch.arm(read)
|
|
|
|
@staticmethod
|
|
def arm_cooldowns(ctx: "RoutingContext") -> None:
|
|
"""The cooldown read routing takes later, declared on the request's native pre-call batch now."""
|
|
routing_prefetch.arm_cooldowns(ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids"))
|
|
|
|
@staticmethod
|
|
def armed() -> bool:
|
|
return routing_prefetch.armed()
|
|
|
|
|
|
_active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar(
|
|
"routing_read_batch", default=None
|
|
)
|
|
|
|
|
|
class RoutingReadBatch:
|
|
def __init__(self, usage_selector: LowestTPMLoggingHandler_v2 | None) -> None:
|
|
self.usage_selector: Final = usage_selector
|
|
self.prefetched_usage: PrefetchedUsage | None = None
|
|
|
|
@staticmethod
|
|
@contextmanager
|
|
def scoped(batch: "RoutingReadBatch | None") -> Iterator[None]:
|
|
token: Final = _active_routing_read_batch.set(batch)
|
|
try:
|
|
yield
|
|
finally:
|
|
_active_routing_read_batch.reset(token)
|
|
|
|
@staticmethod
|
|
def active() -> "RoutingReadBatch | None":
|
|
return _active_routing_read_batch.get()
|
|
|
|
@staticmethod
|
|
def for_strategy(strategy: str | None, selector: object) -> "RoutingReadBatch | None":
|
|
if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2):
|
|
return RoutingReadBatch(usage_selector=selector)
|
|
return RoutingReadBatch(usage_selector=None) if RoutingPrefetch.armed() else None
|
|
|
|
@renamed_router_param("litellm_router_instance")
|
|
async def async_get_cooldown_deployments(
|
|
self, ctx: "RoutingContext", healthy_deployments: Sequence[object], parent_otel_span: "Span | None"
|
|
) -> list[str]:
|
|
cooldowns, _ = await asyncio.gather(
|
|
routing_prefetch.read_cooldowns(
|
|
ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids"), parent_otel_span
|
|
),
|
|
self.read_usage(healthy_deployments, parent_otel_span),
|
|
)
|
|
verbose_router_logger.debug("retrieve cooldown models: %s", cooldowns)
|
|
return [model_id for model_id, _ in cooldowns]
|
|
|
|
async def read_usage(self, deployments: Sequence[object], parent_otel_span: "Span | None") -> None:
|
|
"""The selector's tpm/rpm counters for ``deployments``, kept as ``prefetched_usage``."""
|
|
usage_read: Final = _selector_read(self.usage_selector, deployments)
|
|
if usage_read is None:
|
|
return
|
|
values: Final = await routing_prefetch.read(usage_read, parent_otel_span)
|
|
self.prefetched_usage = PrefetchedUsage(
|
|
keys=frozenset(usage_read.keys),
|
|
values=None if values is None else dict(zip(usage_read.keys, values)),
|
|
)
|