scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

refactor(router): count io-token itpm and otpm on the native router_usage store

Browse source 
The io-token check reserves, reconciles and refunds its global_router itpm and
otpm counters through the native router_usage store, so ModelRateLimitingCheck
keeps no DualCache. A reservation stashes its counter's fields in the request
metadata instead of a Redis key, so nothing parses a key back into ids. The
Redis keys and values are unchanged, pinned by contract tests on both sides.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-10 16:45:58 +09:00
parent 68162dce66
commit b5a7064faf
24 changed files with 678 additions and 656 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs
View file

@ -301,6 +301,18 @@ impl NativeRouterUsage {
})
}
fn shared_or_local<'py>(
&self,
py: Python<'py>,
counter: PyCounter,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_async(py, "async_get_cache", kind_error, async move {
Ok(store.shared_or_local(&counter).await?.map(PyNumber))
})
}
fn current_sync(
&self,
py: Python<'_>,

16
litellm-rust/crates/python-bridge/src/router/mod.rs
View file

@ -385,22 +385,6 @@ impl RouterCoreHandle {
.collect()
}
/// The key deployment `id` serving litellm `model` counts `metric` (`tpm`, `rpm`, `itpm` or
/// `otpm`) under in `minute` (`HH-MM`), as the usage store spells it.
fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult<String> {
let kind: Kind = metric
.parse()
.map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?;
Ok(Counter {
scope: Scope::GlobalRouter,
deployment_id: id.to_owned(),
model: Some(model.to_owned()),
window: minute.to_owned(),
kind,
}
.key())
}
/// Keep the providers `other`'s build resolved for the deployments this router serves too.
fn adopt_providers(&self, other: &Self) {
self.router.adopt_providers(&other.router);

71
litellm-rust/crates/redis-store/src/router_usage/mod.rs
View file

@ -219,28 +219,21 @@ impl RouterUsage {
if local.is_some_and(|local| local.as_f64() >= limit.as_f64()) {
return Ok(local);
}
let Some(client) = self.usage.redis().client() else {
return Ok(local);
};
let key = counter.key();
match client.send(commands::get::<Vec<u8>>(&key)).await {
Err(RedisError::BreakerOpen) => Ok(local),
Err(error) => {
redis_failure!(
&error,
"litellm.caching.caching: get() - Got exception from REDIS"
);
Ok(None)
}
Ok(raw) => Ok(raw
.as_deref()
.map(codec::decode)
.transpose()
.unwrap_or_else(|error| {
tracing::warn!(%error, key, "a usage counter did not decode");
None
})),
}
Ok(match self.shared(counter).await {
Shared::Unasked => local,
Shared::Failed => None,
Shared::Read(count) => count,
})
}
/// The io-token check's usage: what every pod counted in Redis, else this router's memory
/// count when Redis has none, is absent, refuses or fails.
pub async fn shared_or_local(&self, counter: &Counter) -> Result<Option<Number>, Error> {
let local = self.get_local(counter)?;
Ok(match self.shared(counter).await {
Shared::Read(Some(count)) => Some(count),
Shared::Unasked | Shared::Failed | Shared::Read(None) => local,
})
}
pub fn flush_local_counters(&self) -> Result<(), Error> {
@ -252,6 +245,40 @@ impl RouterUsage {
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
Ok(self.usage.memory().entries()?)
}
/// `redis_cache.get_cache` of one counter, which `RedisCache` answers with `None` on a failure
/// and its breaker refuses.
async fn shared(&self, counter: &Counter) -> Shared {
let Some(client) = self.usage.redis().client() else {
return Shared::Unasked;
};
let key = counter.key();
match client.send(commands::get::<Vec<u8>>(&key)).await {
Err(RedisError::BreakerOpen) => Shared::Unasked,
Err(error) => {
redis_failure!(
&error,
"litellm.caching.caching: get() - Got exception from REDIS"
);
Shared::Failed
}
Ok(raw) => match raw.as_deref().map(codec::decode).transpose() {
Ok(count) => Shared::Read(count),
Err(error) => {
tracing::warn!(%error, key, "a usage counter did not decode");
Shared::Failed
}
},
}
}
}
/// What one counter's Redis read came back with.
enum Shared {
/// No Redis, or its breaker refused the read.
Unasked,
Failed,
Read(Option<Number>),
}
fn keys(counters: &[Counter]) -> Vec<String> {

61
litellm-rust/crates/redis-store/tests/router_usage.rs
View file

@ -658,3 +658,64 @@ async fn a_failed_deployment_pipeline_keeps_the_memory_counts_as_integers() {
[Some(Number::Int(5)), Some(Number::Int(1))]
);
}
// The io-token check stashes a counter's fields and the store spells the key, so the key must stay
// main's RouterCacheEnum.ITPM / OTPM spelling
#[rstest]
#[case::itpm(Kind::Itpm, "global_router:dep-1:bedrock/claude:itpm:12-34")]
#[case::otpm(Kind::Otpm, "global_router:dep-1:bedrock/claude:otpm:12-34")]
#[tokio::test]
async fn io_token_counters_increment_under_mains_keys(#[case] kind: Kind, #[case] key: &str) {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let counter = Counter {
scope: Scope::GlobalRouter,
deployment_id: "dep-1".to_owned(),
model: Some("bedrock/claude".to_owned()),
window: "12-34".to_owned(),
kind,
};
store
.increment(&counter, Number::Int(8), COUNTER_TTL)
.await
.unwrap();
let written: Option<String> = raw.send(commands::get(key)).await.unwrap();
assert_eq!(written, Some("8".to_owned()));
}
#[rstest]
#[tokio::test]
async fn shared_usage_prefers_what_every_pod_counted_and_falls_back_to_memory() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let [rpm, tpd] = deployment_usage();
for case in [&rpm, &tpd] {
store
.increment_local(&increment(&case.counter, Number::Int(2), COUNTER_TTL))
.unwrap();
}
raw.send(commands::incr_by_float(
rpm.key,
Amount::Float(9.0),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
let counted = store.shared_or_local(&rpm.counter).await.unwrap();
let absent = store.shared_or_local(&tpd.counter).await.unwrap();
deployment.stop().await.unwrap();
let failed = store.shared_or_local(&rpm.counter).await.unwrap();
assert_eq!(counted, Some(Number::Int(9)));
assert_eq!(
absent,
Some(Number::Int(2)),
"Redis has no count, so memory answers"
);
assert_eq!(failed, Some(Number::Int(2)));
}

6
litellm/litellm_core_utils/internal_call_metadata.py
View file

@ -44,9 +44,9 @@ BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost"
ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved"
OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved"
ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key"
OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key"
IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY})
ITPM_COUNTER_KEY: Final = "_litellm_itpm_counter"
OTPM_COUNTER_KEY: Final = "_litellm_otpm_counter"
IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_COUNTER_KEY, OTPM_COUNTER_KEY})
MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups"
"""Where auth records the model access groups that authorized the request, for the spend writer.

9
litellm/redis/stores/router_usage/__init__.py
View file

@ -66,6 +66,7 @@ class RouterUsage(Protocol):
def set_local(self, counter: UsageCounter, value: float) -> None: ...
async def current(self, counter: UsageCounter, limit: float) -> float | None: ...
def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ...
async def shared_or_local(self, counter: UsageCounter) -> float | None: ...
def snapshot(self) -> Mapping[str, float]: ...
@ -197,5 +198,13 @@ class RouterUsageStore:
def current_sync(self, counter: UsageCounter, limit: float) -> float | None:
return self._handle.current_sync(counter, limit)
@with_service_target(ROUTER_USAGE_TARGET)
async def shared_or_local(self, counter: UsageCounter) -> float | None:
"""What every pod counted in Redis, else this router's memory count when Redis has none, is absent, refuses
or fails."""
if not self.has_redis:
return self._handle.get_local(counter)
return await self._handle.shared_or_local(counter)
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()

1
litellm/redis/stores/router_usage/handle.py
View file

@ -42,5 +42,6 @@ class RouterUsageHandle(Protocol):
def set_local(self, counter: UsageCounter, value: float) -> None: ...
def current(self, counter: UsageCounter, limit: float) -> Awaitable[float | None]: ...
def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ...
def shared_or_local(self, counter: UsageCounter) -> Awaitable[float | None]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...

6
litellm/router_rust/core.py
View file

@ -1290,6 +1290,10 @@ class _EndpointPort:
def routing_cache(self) -> "DualCache":
return self._core._cache # pyright: ignore[reportPrivateUsage] # see the class docstring
@property
def router_usage(self) -> RouterUsage:
return self._core._usage # pyright: ignore[reportPrivateUsage] # see the class docstring
@property
def selectors(self) -> StrategySelectors:
return self._core._selectors # pyright: ignore[reportPrivateUsage] # see the class docstring
@ -5100,7 +5104,7 @@ class _Driver:
# logging object (shared across attempts, it logs only the first failure).
with contextlib.suppress(Exception):
await async_io_token_refund_failure(
self._core._cache, # pyright: ignore[reportPrivateUsage] # the driver serves its own core
self._core._usage, # pyright: ignore[reportPrivateUsage] # the driver serves its own core
call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict
parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs),
)

27
litellm/router_rust/invoke/endpoints.py
View file

@ -56,6 +56,7 @@ from litellm.router_utils.endpoint_eligibility import endpoint_for_call, endpoin
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.redis.stores.router_usage import RouterUsage
from litellm.router_api import RouterRegistry
from litellm.router_rust.core import RustCore
from litellm.router_rust.steps.selection import StrategySelectors
@ -93,6 +94,9 @@ class EndpointCore(Protocol):
@property
def routing_cache(self) -> "DualCache": ...
@property
def router_usage(self) -> "RouterUsage": ...
@property
def selectors(self) -> "StrategySelectors": ...
@ -851,7 +855,6 @@ def attempt_call(
named: Final = {name: kwargs.pop(name) for name in plan.named if name in kwargs}
data: Final = as_dict(deployment.get("litellm_params"))
model_name: Final = str(data.get("model"))
cache: Final = core.routing_cache
caching: Final = shaping.cache_responses
# A fallback hop routes the target group in a copy of the metadata (`run_async_fallback`).
if fallback_level > 0:
@ -868,7 +871,9 @@ def attempt_call(
start_silent_experiments(core, plan, silent_model, named.get("messages"), kwargs)
if "messages" in named:
kwargs.setdefault("messages", named["messages"])
update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, splits_stream_timeout=True)
update_kwargs_with_deployment(
shaping, deployment, kwargs, usage=core.router_usage, splits_stream_timeout=True
)
kwargs.pop("silent_model", None)
final = without_router_only_kwargs( # rebind-ok: rebound as the Python router rebinds it
{
@ -882,7 +887,7 @@ def attempt_call(
return AttemptCall(plan.host_op, final, kwargs, str(params.get("model")), group)
case "text" | "embedding" | "image" | "transcription" | "rerank" | "adapter":
update_kwargs_with_deployment(
shaping, deployment, kwargs, cache=cache, splits_stream_timeout=plan.splits_stream_timeout
shaping, deployment, kwargs, usage=core.router_usage, splits_stream_timeout=plan.splits_stream_timeout
)
refused: Final = embedding_batch_refusal(plan, deployment, named.get("input"), kwargs, requested)
if refused is not None:
@ -896,7 +901,7 @@ def attempt_call(
}
return AttemptCall(plan.host_op, without_router_only_kwargs(final), kwargs, model_name, group)
case "speech":
update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache)
update_kwargs_with_deployment(shaping, deployment, kwargs, usage=core.router_usage)
voice = kwargs.pop("voice", None) # rebind-ok: rebound as the Python router rebinds it
final = { # rebind-ok: rebound as the Python router rebinds it
**data,
@ -910,7 +915,9 @@ def attempt_call(
if plan.kind == "cancel_batch":
data.update(_resolved_credentials(core, deployment, requested))
data.pop("litellm_credential_name", None)
update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, function_name=plan.function_name)
update_kwargs_with_deployment(
shaping, deployment, kwargs, usage=core.router_usage, function_name=plan.function_name
)
final = { # rebind-ok: rebound as the Python router rebinds it
**data,
"custom_llm_provider": _provider(core, data),
@ -926,7 +933,7 @@ def attempt_call(
shaping,
deployment,
kwargs,
cache=cache,
usage=core.router_usage,
function_name="generic_api_call",
splits_stream_timeout=invoked_handler_name(core, plan) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
)
@ -942,7 +949,7 @@ def attempt_call(
}
return AttemptCall(plan.host_op, final, kwargs, model_name, group)
case "moderation":
update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache)
update_kwargs_with_deployment(shaping, deployment, kwargs, usage=core.router_usage)
kwargs["model"] = model_name
kwargs.update(data)
return AttemptCall(plan.host_op, kwargs, kwargs, model_name, group)
@ -973,7 +980,7 @@ def _generic_attempt(
core.shaping,
deployment,
kwargs,
cache=core.routing_cache,
usage=core.router_usage,
function_name="_ageneric_api_call_with_fallbacks",
splits_stream_timeout=handler in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
opens_responses_socket=socket,
@ -2205,7 +2212,7 @@ class RustEndpoints:
shaping,
deployment,
new_kwargs, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
cache=self._core.routing_cache,
usage=self._core.router_usage,
function_name="aretrieve_batch",
)
bucket: Final = metadata_bucket(new_kwargs, metadata_variable_name) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
@ -2289,7 +2296,7 @@ class RustEndpoints:
self._core.shaping,
deployment,
kwargs_copy, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
cache=self._core.routing_cache,
usage=self._core.router_usage,
function_name="acreate_file",
)
data: Final = as_dict(deployment.get("litellm_params"))

2
litellm/router_rust/invoke/prompt_management.py
View file

@ -76,7 +76,7 @@ async def prompt_management_acompletion(
core.shaping,
deployment,
kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs
cache=core.routing_cache,
usage=core.router_usage,
splits_stream_timeout=True,
)
params: Final = as_dict(deployment.get("litellm_params"))

8
litellm/router_rust/invoke/shaping.py
View file

@ -52,7 +52,7 @@ from litellm.types.router import Deployment, GenericLiteLLMParams, LiteLLM_Param
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.redis.stores.router_usage import RouterUsage
from litellm.router_python_host import PythonHost
# The SDK-client and max_parallel_requests reads of the routing cache: process-local, so the family
@ -403,7 +403,7 @@ def update_kwargs_with_deployment(
deployment: Mapping[str, object],
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
*,
cache: "DualCache | None",
usage: "RouterUsage | None",
function_name: str | None = None,
splits_stream_timeout: bool = False,
opens_responses_socket: bool = False,
@ -466,8 +466,8 @@ def update_kwargs_with_deployment(
if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata:
kwargs[other_bucket] = {**as_dict(other_bucket_metadata), KUBERNETES_POD_ROUTING_KEY: routing} # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
if cache is not None:
refund_stale_reservation_before_retry(cache, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict
if usage is not None:
refund_stale_reservation_before_retry(usage, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict
set_io_token_rate_limit_request_kwargs(
kwargs, # pyright: ignore[reportArgumentType] # the request's kwargs dict
store_in_context=deployment_has_io_token_limits(deployment), # pyright: ignore[reportArgumentType] # a deployment record

2
litellm/router_rust/invoke/websocket.py
View file

@ -64,7 +64,7 @@ async def admit_deployment_turn(
io_token_admission: Final[Mapping[str, object]] = (
MappingProxyType({})
if rate_limit_check is None
else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict
else await async_admit_io_token_turn(rate_limit_check.usage, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict
)
budget_limiter: Final[object] = _budget_limiter(core)
if budget_limiter is not None:

2
litellm/router_rust/steps/callbacks.py
View file

@ -306,7 +306,7 @@ class OptionalChecks:
)
self.router_budget_logger = callback
elif name == "enforce_model_rate_limits":
callback = ModelRateLimitingCheck(usage=self._usage, dual_cache=self._cache)
callback = ModelRateLimitingCheck(usage=self._usage)
if callback is None:
continue
if self.optional_callbacks is not None and any(

2
litellm/router_rust/surface/legacy_members.py
View file

@ -421,7 +421,7 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks):
except Exception:
with contextlib.suppress(Exception):
await async_io_token_refund_failure(
self._cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
self._usage, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
)
raise
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)

46
litellm/router_rust/surface/model_info.py
View file

@ -7,8 +7,8 @@ host registry.
router-core holds the rules: a deployment's model info (``get_router_model_info``,
``get_deployment_model_info``: ``router-core/src/model_info.rs``), a model group's info and the
configured-value reads (``_set_model_group_info``, ``get_model_listing_info``, the configured token
limits, mode, display name and service tiers: ``router-core/src/group_info.rs``), and the usage
counters' keys (``RouterCoreHandle.usage_key``). What lives in this process is answered here as a
limits, mode, display name and service tiers: ``router-core/src/group_info.rs``); the usage
counters' keys are the ``router_usage`` store's. What lives in this process is answered here as a
query asks for it, at query time: the cost map (``get_model_info``, ``litellm.model_cost``),
provider resolution, the wildcard pattern routes, what the background refresh discovered, and a
group deployment's validated params, supported params, endpoints, fast mode and reasoning efforts.
@ -33,7 +33,6 @@ from pydantic import BaseModel, TypeAdapter
from typing_extensions import overload
import litellm
from litellm._internal_context import service_target
from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.add_retry_fallback_headers import (
apply_quality_router_decision_headers,
@ -54,7 +53,7 @@ from litellm.litellm_core_utils.reasoning_effort_capability import (
from litellm.litellm_core_utils.routing_common_utils import team_may_use_deployment
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter
from litellm.redis.stores.router_usage import RouterUsage, UsageCounter
from litellm.router_python_host import python_host
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import build_io_token_rate_limit_headers
from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata
@ -797,30 +796,31 @@ class HostModelInfo:
"""
Returns current ITPM/OTPM usage for a model group (sum across deployments).
"""
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
itpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged
otpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged
current_minute: Final = get_utc_datetime().strftime("%H-%M")
model_list: Final = self._core_get_model_list(model_name=model_group)
if model_list is None:
return None, None
for model in model_list:
model_id: str | None = model.get("model_info", {}).get("id")
litellm_model: str | None = model["litellm_params"].get("model")
if model_id is None or litellm_model is None:
continue
itpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "itpm", current_minute))
otpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "otpm", current_minute))
named: Final[tuple[tuple[object, object], ...]] = tuple(
(model.get("model_info", {}).get("id"), model["litellm_params"].get("model")) for model in model_list
)
deployments: Final = tuple(
(f"{model_id}", f"{litellm_model}")
for model_id, litellm_model in named
if model_id is not None and litellm_model is not None
)
itpm_counters: Final = tuple(
UsageCounter("global_router", model_id, litellm_model, current_minute, "itpm")
for model_id, litellm_model in deployments
)
otpm_counters: Final = tuple(
UsageCounter("global_router", model_id, litellm_model, current_minute, "otpm")
for model_id, litellm_model in deployments
)
combined_values: Final = await self._core_usage_store.get_many((*itpm_counters, *otpm_counters))
with service_target(ROUTER_USAGE_TARGET):
combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys)
if combined_values is None:
return None, None
itpm_values: Final = combined_values[: len(itpm_keys)]
otpm_values: Final = combined_values[len(itpm_keys) :]
itpm_values: Final = combined_values[: len(itpm_counters)]
otpm_values: Final = combined_values[len(itpm_counters) :]
total_itpm: int | None = None # rebind-ok: ported from the Python core unchanged
for value in itpm_values:

277
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py
View file

@ -14,36 +14,35 @@ from __future__ import annotations
import contextlib
import contextvars
from collections.abc import Mapping, MutableMapping
from dataclasses import dataclass
from dataclasses import asdict, dataclass, replace
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import httpx
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm import token_counter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.core_helpers import coerce_token_limit
from litellm.litellm_core_utils.internal_call_metadata import (
IO_TOKEN_RESERVATION_KEYS,
ITPM_CACHE_KEY,
ITPM_COUNTER_KEY,
ITPM_RESERVED_KEY,
OTPM_CACHE_KEY,
OTPM_COUNTER_KEY,
OTPM_RESERVED_KEY,
)
from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET
from litellm.types.router import RouterCacheEnum, RouterErrors
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter
from litellm.types.router import RouterErrors
from litellm.utils import get_utc_datetime
_MODEL_INFO_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_STASHED_COUNTER: Final[TypeAdapter[UsageCounter]] = TypeAdapter(UsageCounter)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -144,16 +143,17 @@ def _without_io_token_limit_keys(section: object) -> object:
return {key: value for key, value in section.items() if key not in _IO_TOKEN_LIMIT_KEYS}
def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | None:
model_id: Final = deployment.get("model_info", {}).get("id")
deployment_name: Final = deployment.get("litellm_params", {}).get("model")
def _get_counters(deployment: dict, current_minute: str) -> tuple[UsageCounter, UsageCounter] | None:
model_id: Final[object] = deployment.get("model_info", {}).get("id")
deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model")
# Without both a deployment id and model name the key would collapse to a
# shared "None:None" bucket across misconfigured deployments, so bail out.
if model_id is None or deployment_name is None:
return None
itpm_key = RouterCacheEnum.ITPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute)
otpm_key = RouterCacheEnum.OTPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute)
return itpm_key, otpm_key
return (
UsageCounter("global_router", f"{model_id}", f"{deployment_name}", current_minute, "itpm"),
UsageCounter("global_router", f"{model_id}", f"{deployment_name}", current_minute, "otpm"),
)
def _served_model(deployment: object) -> str:
@ -354,14 +354,15 @@ def _resolve_reconcile_usage_tokens(
def _reservation(
itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None
itpm_reserved: int, otpm_reserved: int, itpm: UsageCounter | None, otpm: UsageCounter | None
) -> Mapping[str, object]:
"""The reservation as plain fields, which survive the copies logging makes of the request metadata."""
return MappingProxyType(
{
ITPM_RESERVED_KEY: itpm_reserved,
OTPM_RESERVED_KEY: otpm_reserved,
ITPM_CACHE_KEY: itpm_cache_key,
OTPM_CACHE_KEY: otpm_cache_key,
ITPM_COUNTER_KEY: None if itpm is None else asdict(itpm),
OTPM_COUNTER_KEY: None if otpm is None else asdict(otpm),
}
)
@ -372,29 +373,40 @@ def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, res
metadata.update(reservation)
def _key_for_this_minute(reserved_key: str) -> str:
counter, _, _ = reserved_key.rpartition(":")
return f"{counter}:{get_utc_datetime().strftime('%H-%M')}"
def _this_minute(counter: UsageCounter) -> UsageCounter:
return replace(counter, window=get_utc_datetime().strftime("%H-%M"))
_Reservation = tuple[int, int, UsageCounter | None, UsageCounter | None]
def _reconcile_increments(
reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int
) -> tuple[tuple[str, int], ...]:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved))
reservation: _Reservation, billable_input: int, completion_tokens: int
) -> tuple[tuple[UsageCounter, int], ...]:
itpm_reserved, otpm_reserved, itpm, otpm = reservation
deltas: Final = ((itpm, billable_input - itpm_reserved), (otpm, completion_tokens - otpm_reserved))
return tuple(
(key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta
(counter if delta < 0 else _this_minute(counter), delta)
for counter, delta in deltas
if counter is not None and delta
)
def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]:
itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY)
otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY)
def _stashed_counter(value: object) -> UsageCounter | None:
if value is None:
return None
try:
return _STASHED_COUNTER.validate_python(value)
except ValidationError:
return None
def _extract_reservation(reservation: Mapping[str, object]) -> _Reservation:
return (
int(reservation.get(ITPM_RESERVED_KEY, 0) or 0),
int(reservation.get(OTPM_RESERVED_KEY, 0) or 0),
itpm_cache_key if isinstance(itpm_cache_key, str) else None,
otpm_cache_key if isinstance(otpm_cache_key, str) else None,
int(reservation.get(ITPM_RESERVED_KEY, 0) or 0), # pyright: ignore[reportArgumentType] # the int the router stashed
int(reservation.get(OTPM_RESERVED_KEY, 0) or 0), # pyright: ignore[reportArgumentType] # see above
_stashed_counter(reservation.get(ITPM_COUNTER_KEY)),
_stashed_counter(reservation.get(OTPM_COUNTER_KEY)),
)
@ -421,7 +433,7 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object,
return top_level + from_params + from_logging_object
def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> tuple[int, int, str | None, str | None]:
def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> _Reservation:
for channel_dict in _reservation_channels(kwargs):
if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict:
return _extract_reservation(channel_dict)
@ -476,8 +488,8 @@ def _rate_limit_error(limit_label: str, limit: int, current: float) -> litellm.R
@with_service_target(ROUTER_USAGE_TARGET)
def _sync_increment_with_rollback(
dual_cache: DualCache,
key: str,
usage: RouterUsage,
counter: UsageCounter,
value: int,
limit: int | None,
*,
@ -485,24 +497,16 @@ def _sync_increment_with_rollback(
) -> None:
if value <= 0 or limit is None:
return
current: Final = dual_cache.increment_cache(
key=key,
value=value,
ttl=RoutingArgsTTL,
)
current: Final = usage.increment_sync(counter, value, RoutingArgsTTL)
if current is not None and current > limit:
dual_cache.increment_cache(
key=key,
value=-value,
ttl=RoutingArgsTTL,
)
usage.increment_sync(counter, -value, RoutingArgsTTL)
raise _rate_limit_error(limit_label, limit, current)
@with_service_target(ROUTER_USAGE_TARGET)
async def _increment_with_rollback(
dual_cache: DualCache,
key: str,
usage: RouterUsage,
counter: UsageCounter,
value: int,
limit: int | None,
*,
@ -511,25 +515,15 @@ async def _increment_with_rollback(
) -> None:
if value <= 0 or limit is None:
return
current: Final = await dual_cache.async_increment_cache(
key=key,
value=value,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
current: Final = await usage.increment(counter, value, RoutingArgsTTL)
if current is not None and current > limit:
await dual_cache.async_increment_cache(
key=key,
value=-value,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
await usage.increment(counter, -value, RoutingArgsTTL)
raise _rate_limit_error(limit_label, limit, current)
@with_service_target(ROUTER_USAGE_TARGET)
def io_token_pre_call_check(
dual_cache: DualCache,
usage: RouterUsage,
deployment: dict,
) -> dict | None:
itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment)
@ -546,10 +540,10 @@ def io_token_pre_call_check(
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
cache_keys: Final = _get_cache_keys(deployment, current_minute)
if cache_keys is None:
counters: Final = _get_counters(deployment, current_minute)
if counters is None:
return deployment
itpm_key, otpm_key = cache_keys
itpm_counter, otpm_counter = counters
itpm_reserved = 0
otpm_reserved = 0
@ -557,8 +551,8 @@ def io_token_pre_call_check(
if itpm_limit is not None:
itpm_reserved = _reservation_value(estimated_input, itpm_limit)
_sync_increment_with_rollback(
dual_cache,
itpm_key,
usage,
itpm_counter,
itpm_reserved,
itpm_limit,
limit_label="ITPM",
@ -568,19 +562,15 @@ def io_token_pre_call_check(
otpm_reserved = otpm_to_reserve
try:
_sync_increment_with_rollback(
dual_cache,
otpm_key,
usage,
otpm_counter,
otpm_reserved,
otpm_limit,
limit_label="OTPM",
)
except Exception:
if itpm_reserved > 0:
dual_cache.increment_cache(
key=itpm_key,
value=-itpm_reserved,
ttl=RoutingArgsTTL,
)
usage.increment_sync(itpm_counter, -itpm_reserved, RoutingArgsTTL)
raise
_stash_reservation_in_metadata(
@ -588,8 +578,8 @@ def io_token_pre_call_check(
_reservation(
itpm_reserved,
otpm_reserved,
itpm_key if itpm_limit is not None else None,
otpm_key if otpm_limit is not None else None,
itpm_counter if itpm_limit is not None else None,
otpm_counter if otpm_limit is not None else None,
),
)
return deployment
@ -597,7 +587,7 @@ def io_token_pre_call_check(
@with_service_target(ROUTER_USAGE_TARGET)
async def async_io_token_pre_call_check(
dual_cache: DualCache,
usage: RouterUsage,
deployment: dict,
parent_otel_span: Span | None = None,
) -> dict | None:
@ -610,7 +600,7 @@ async def async_io_token_pre_call_check(
if request is not None and request.opens_responses_socket:
_stash_reservation_in_metadata(
request,
await async_admit_io_token_turn(dual_cache, deployment, get_utc_datetime(), parent_otel_span),
await async_admit_io_token_turn(usage, deployment, get_utc_datetime(), parent_otel_span),
)
return deployment
_model, served_model = _tokenizer_and_served_models(deployment)
@ -621,10 +611,10 @@ async def async_io_token_pre_call_check(
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M")
cache_keys: Final = _get_cache_keys(deployment, current_minute)
if cache_keys is None:
counters: Final = _get_counters(deployment, current_minute)
if counters is None:
return deployment
itpm_key, otpm_key = cache_keys
itpm_counter, otpm_counter = counters
itpm_reserved = 0
otpm_reserved = 0
@ -632,8 +622,8 @@ async def async_io_token_pre_call_check(
if itpm_limit is not None:
itpm_reserved = _reservation_value(estimated_input, itpm_limit)
await _increment_with_rollback(
dual_cache,
itpm_key,
usage,
itpm_counter,
itpm_reserved,
itpm_limit,
parent_otel_span=parent_otel_span,
@ -644,8 +634,8 @@ async def async_io_token_pre_call_check(
otpm_reserved = otpm_to_reserve
try:
await _increment_with_rollback(
dual_cache,
otpm_key,
usage,
otpm_counter,
otpm_reserved,
otpm_limit,
parent_otel_span=parent_otel_span,
@ -656,12 +646,7 @@ async def async_io_token_pre_call_check(
# release the ITPM reservation already made, or it stays inflated
# until the TTL expires.
if itpm_reserved > 0:
await dual_cache.async_increment_cache(
key=itpm_key,
value=-itpm_reserved,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
await usage.increment(itpm_counter, -itpm_reserved, RoutingArgsTTL)
raise
_stash_reservation_in_metadata(
@ -669,42 +654,26 @@ async def async_io_token_pre_call_check(
_reservation(
itpm_reserved,
otpm_reserved,
itpm_key if itpm_limit is not None else None,
otpm_key if otpm_limit is not None else None,
itpm_counter if itpm_limit is not None else None,
otpm_counter if otpm_limit is not None else None,
),
)
return deployment
async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None:
if redis_cache is None:
return None
with contextlib.suppress(RedisCircuitBreakerOpenError):
return await redis_cache.async_get_cache(key=key, parent_otel_span=parent_otel_span)
return None
@with_service_target(ROUTER_USAGE_TARGET)
async def _async_shared_usage(dual_cache: DualCache, key: str, parent_otel_span: Span | None) -> int | None:
shared: Final = await _async_redis_usage(dual_cache.redis_cache, key, parent_otel_span)
if shared is not None:
return shared
return await dual_cache.async_get_cache(key=key, local_only=True)
async def _raise_if_at_limit(
dual_cache: DualCache, key: str, limit: int | None, limit_label: str, parent_otel_span: Span | None
usage: RouterUsage, counter: UsageCounter, limit: int | None, limit_label: str, parent_otel_span: Span | None
) -> None:
if limit is None:
return
usage: Final = await _async_shared_usage(dual_cache, key, parent_otel_span)
if usage is not None and usage >= limit:
raise _rate_limit_error(limit_label, limit, usage)
count: Final = await usage.shared_or_local(counter)
if count is not None and count >= limit:
raise _rate_limit_error(limit_label, limit, count)
@with_service_target(ROUTER_USAGE_TARGET)
async def async_admit_io_token_turn(
dual_cache: DualCache,
usage: RouterUsage,
deployment: Mapping[str, object],
now: datetime,
parent_otel_span: Span | None = None,
@ -713,36 +682,36 @@ async def async_admit_io_token_turn(
itpm_limit, otpm_limit = get_deployment_io_token_limits(configured)
if itpm_limit is None and otpm_limit is None:
return MappingProxyType({})
cache_keys: Final = _get_cache_keys(configured, now.strftime("%H-%M"))
if cache_keys is None:
counters: Final = _get_counters(configured, now.strftime("%H-%M"))
if counters is None:
return MappingProxyType({})
itpm_key, otpm_key = cache_keys
await _raise_if_at_limit(dual_cache, itpm_key, itpm_limit, "ITPM", parent_otel_span)
await _raise_if_at_limit(dual_cache, otpm_key, otpm_limit, "OTPM", parent_otel_span)
itpm_counter, otpm_counter = counters
await _raise_if_at_limit(usage, itpm_counter, itpm_limit, "ITPM", parent_otel_span)
await _raise_if_at_limit(usage, otpm_counter, otpm_limit, "OTPM", parent_otel_span)
return _reservation(
0, 0, itpm_key if itpm_limit is not None else None, otpm_key if otpm_limit is not None else None
0, 0, itpm_counter if itpm_limit is not None else None, otpm_counter if otpm_limit is not None else None
)
@with_service_target(ROUTER_USAGE_TARGET)
def io_token_reconcile_success(
dual_cache: DualCache,
usage: RouterUsage,
kwargs: Mapping[str, object] | None,
response_obj: object,
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation
if itpm_counter is None and otpm_counter is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
try:
if usage_resolved:
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
for counter, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
usage.increment_sync(counter, delta, RoutingArgsTTL)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
@ -764,7 +733,7 @@ def io_token_reconcile_success(
@with_service_target(ROUTER_USAGE_TARGET)
async def async_io_token_reconcile_success(
dual_cache: DualCache,
usage: RouterUsage,
kwargs: Mapping[str, object] | None,
response_obj: object,
*,
@ -773,8 +742,8 @@ async def async_io_token_reconcile_success(
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation
if itpm_counter is None and otpm_counter is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
@ -785,10 +754,8 @@ async def async_io_token_reconcile_success(
# event can't re-process it.
try:
if usage_resolved:
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
await dual_cache.async_increment_cache(
key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span
)
for counter, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
await usage.increment(counter, delta, RoutingArgsTTL)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
@ -810,30 +777,22 @@ async def async_io_token_reconcile_success(
@with_service_target(ROUTER_USAGE_TARGET)
def io_token_refund_failure(
dual_cache: DualCache,
usage: RouterUsage,
kwargs: Mapping[str, object] | None,
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = _read_reservation_from_kwargs(request_kwargs)
if itpm_counter is None and otpm_counter is None:
return
if itpm_key is not None and itpm_reserved > 0:
dual_cache.increment_cache(
key=itpm_key,
value=-itpm_reserved,
ttl=RoutingArgsTTL,
)
if otpm_key is not None and otpm_reserved > 0:
dual_cache.increment_cache(
key=otpm_key,
value=-otpm_reserved,
ttl=RoutingArgsTTL,
)
if itpm_counter is not None and itpm_reserved > 0:
usage.increment_sync(itpm_counter, -itpm_reserved, RoutingArgsTTL)
if otpm_counter is not None and otpm_reserved > 0:
usage.increment_sync(otpm_counter, -otpm_reserved, RoutingArgsTTL)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)
def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping[str, object] | None) -> None:
def refund_stale_reservation_before_retry(usage: RouterUsage, kwargs: Mapping[str, object] | None) -> None:
"""
Synchronously refund and clear any reservation a previous deployment
attempt stashed in ``kwargs``, before it's overwritten for the next
@ -857,34 +816,24 @@ def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping
"""
if not kwargs:
return
io_token_refund_failure(dual_cache, kwargs)
io_token_refund_failure(usage, kwargs)
@with_service_target(ROUTER_USAGE_TARGET)
async def async_io_token_refund_failure(
dual_cache: DualCache,
usage: RouterUsage,
kwargs: Mapping[str, object] | None,
*,
parent_otel_span: Span | None = None,
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = _read_reservation_from_kwargs(request_kwargs)
if itpm_counter is None and otpm_counter is None:
return
if itpm_key is not None and itpm_reserved > 0:
await dual_cache.async_increment_cache(
key=itpm_key,
value=-itpm_reserved,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
if otpm_key is not None and otpm_reserved > 0:
await dual_cache.async_increment_cache(
key=otpm_key,
value=-otpm_reserved,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
if itpm_counter is not None and itpm_reserved > 0:
await usage.increment(itpm_counter, -itpm_reserved, RoutingArgsTTL)
if otpm_counter is not None and otpm_reserved > 0:
await usage.increment(otpm_counter, -otpm_reserved, RoutingArgsTTL)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)

22
litellm/router_utils/pre_call_checks/model_rate_limit_check.py
View file

@ -22,7 +22,6 @@ from pydantic import BaseModel, ConfigDict
import litellm
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, CounterIncrement, RouterUsage, UsageCounter
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
@ -138,11 +137,8 @@ class ModelRateLimitingCheck(CustomLogger):
this check actively enforces those limits across ALL routing strategies.
"""
def __init__(
self, usage: RouterUsage, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime
) -> None:
def __init__(self, usage: RouterUsage, clock: Callable[[], datetime] = get_utc_datetime) -> None:
self.usage = usage
self.dual_cache = dual_cache
self.clock = clock
# model_ids already warned about itpm/otpm + tpm/rpm on the same deployment,
# so the warning is logged once per deployment rather than per request.
@ -200,7 +196,7 @@ class ModelRateLimitingCheck(CustomLogger):
def _refund_io_token_reservation_if_any(self) -> None:
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
if request_kwargs is not None:
io_token_refund_failure(self.dual_cache, request_kwargs)
io_token_refund_failure(self.usage, request_kwargs)
async def _async_refund_io_token_reservation_if_any(
self,
@ -209,7 +205,7 @@ class ModelRateLimitingCheck(CustomLogger):
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
if request_kwargs is not None:
await async_io_token_refund_failure(
self.dual_cache,
self.usage,
request_kwargs,
parent_otel_span=parent_otel_span,
)
@ -280,7 +276,7 @@ class ModelRateLimitingCheck(CustomLogger):
if deployment_has_io_token_limits(deployment):
self._warn_io_token_and_tpm_rpm_coexist_once(deployment)
io_token_pre_call_check(
self.dual_cache,
self.usage,
deployment,
)
io_reservation_made = True
@ -364,7 +360,7 @@ class ModelRateLimitingCheck(CustomLogger):
if deployment_has_io_token_limits(deployment):
self._warn_io_token_and_tpm_rpm_coexist_once(deployment)
await async_io_token_pre_call_check(
self.dual_cache,
self.usage,
deployment,
parent_otel_span=parent_otel_span,
)
@ -453,7 +449,7 @@ class ModelRateLimitingCheck(CustomLogger):
try:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
await async_io_token_reconcile_success(
self.dual_cache,
self.usage,
kwargs,
response_obj,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
@ -500,7 +496,7 @@ class ModelRateLimitingCheck(CustomLogger):
# Never fail the primary logging pipeline over an io-token refund error.
with contextlib.suppress(Exception):
await async_io_token_refund_failure(
self.dual_cache,
self.usage,
kwargs,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
)
@ -513,7 +509,7 @@ class ModelRateLimitingCheck(CustomLogger):
"""
try:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
io_token_reconcile_success(self.dual_cache, kwargs, response_obj)
io_token_reconcile_success(self.usage, kwargs, response_obj)
if standard_logging_object is None:
return
@ -545,6 +541,6 @@ class ModelRateLimitingCheck(CustomLogger):
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
with contextlib.suppress(Exception):
io_token_refund_failure(
self.dual_cache,
self.usage,
kwargs,
)

2
litellm/rust_bridge/_native.pyi
View file

@ -184,6 +184,7 @@ class NativeRouterUsage:
def set_local(self, counter: _UsageCounter, value: float) -> None: ...
def current(self, counter: _UsageCounter, limit: float) -> Future[float | None]: ...
def current_sync(self, counter: _UsageCounter, limit: float) -> float | None: ...
def shared_or_local(self, counter: _UsageCounter) -> Future[float | None]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> dict[str, float]: ...
@ -1089,7 +1090,6 @@ class RouterCoreHandle:
@staticmethod
def configured_service_tiers(model_infos: list[object]) -> list[object]: ...
# The routing-cache key a deployment counts `metric` (tpm, rpm, itpm, otpm) under in `minute`.
def usage_key(self, id: str, model: str, metric: str, minute: str) -> str: ...
def registry_version(self) -> int: ...
# The version the deployments before the newest appends last changed at.
def registry_rewritten(self) -> int: ...

6
tests/unit/redis/stores/router_usage_double.py
View file

@ -100,6 +100,12 @@ class FakeUsageHandle:
def current(self, counter: UsageCounter, limit: float) -> asyncio.Future[float | None]:
return _done(self.current_sync(counter, limit))
def shared_or_local(self, counter: UsageCounter) -> asyncio.Future[float | None]:
self.calls.append(("shared_or_local", (counter,)))
redis: Final = self._reachable()
shared: Final = None if redis is None else redis.get(counter)
return _done(self.local.get(counter) if shared is None else shared)
def _read(self, counters: Sequence[UsageCounter]) -> list[float | None]:
redis: Final = self._reachable()
return [self.local.get(c, None if redis is None else redis.get(c)) for c in counters]

81
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -48,10 +48,11 @@ from litellm.llms.volcengine.responses.transformation import (
VolcEngineResponsesAPIConfig,
)
from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter
from litellm.types.utils import GenericStreamingChunk
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
@ -1701,8 +1702,8 @@ def _guard_recording(bill_refusal: MagicMock):
return guard
_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-34"
_TURNED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-35"
_OPENED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-34", "otpm")
_TURNED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-35", "otpm")
_OPENED_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc)
_TURNED_AT: Final = datetime(2026, 1, 1, 12, 35, 5, tzinfo=timezone.utc)
_WS_DEPLOYMENT: Final = {
@ -1728,12 +1729,12 @@ def _turn_completed() -> str:
)
async def _socket_opened_on(dual_cache: "DualCache") -> Mapping[str, object]:
async def _socket_opened_on(usage: RouterUsageStore) -> Mapping[str, object]:
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn
await dual_cache.async_increment_cache(key=_OPENED_ON, value=_OTHER_USAGE, ttl=60)
await dual_cache.async_increment_cache(key=_TURNED_ON, value=_OTHER_USAGE, ttl=60)
return await async_admit_io_token_turn(dual_cache, _WS_DEPLOYMENT, _OPENED_AT)
usage.increment_local(CounterIncrement(_OPENED_ON, _OTHER_USAGE, 60))
usage.increment_local(CounterIncrement(_TURNED_ON, _OTHER_USAGE, 60))
return await async_admit_io_token_turn(usage, _WS_DEPLOYMENT, _OPENED_AT)
def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> None:
@ -1742,13 +1743,13 @@ def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) ->
)
def _later_turns_admitted_through(dual_cache: "DualCache") -> "ResponsesWsTurnAdmission":
def _later_turns_admitted_through(usage: RouterUsageStore) -> "ResponsesWsTurnAdmission":
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn
return partial(async_admit_io_token_turn, dual_cache, _WS_DEPLOYMENT, _TURNED_AT)
return partial(async_admit_io_token_turn, usage, _WS_DEPLOYMENT, _TURNED_AT)
_ITPM_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:itpm:12-34"
_ITPM_OPENED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-34", "itpm")
_ITPM_DEPLOYMENT: Final = {
"litellm_params": {"model": "openai/gpt-5-mini", "itpm": 40_000},
"model_info": {"id": "ws-dep"},
@ -1775,12 +1776,14 @@ def _turn_completed_reading_cache(cached_tokens: int) -> str:
)
def _rate_limit_check(dual_cache: "DualCache") -> "ModelRateLimitingCheck":
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_usage import RouterUsageStore
def _memory_usage() -> RouterUsageStore:
return RouterUsageStore(connect(None).router_usage())
def _rate_limit_check(usage: RouterUsageStore) -> "ModelRateLimitingCheck":
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
return ModelRateLimitingCheck(usage=RouterUsageStore(connect(None).router_usage()), dual_cache=dual_cache)
return ModelRateLimitingCheck(usage=usage)
def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging":
@ -2012,13 +2015,11 @@ class TestWebSocketTurnAdmission:
) -> None:
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.caching.dual_cache import DualCache
from litellm.responses.streaming_iterator import drain_logging_tasks
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
_turns_finish_at(monkeypatch, _TURNED_AT)
dual_cache: Final = DualCache()
opened_with: Final = await _socket_opened_on(dual_cache)
usage: Final = _memory_usage()
opened_with: Final = await _socket_opened_on(usage)
first, second, third = _turn_frames()
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock()
@ -2029,18 +2030,18 @@ class TestWebSocketTurnAdmission:
websocket=_ClientSocket(second, third),
backend_ws=backend_ws,
request_data={},
logging_obj=_connection_logged_through(_rate_limit_check(dual_cache), opened_with),
logging_obj=_connection_logged_through(_rate_limit_check(usage), opened_with),
first_message=first,
turn_admission=_later_turns_admitted_through(dual_cache),
turn_admission=_later_turns_admitted_through(usage),
)
await handler.client_to_backend()
await handler.backend_to_client()
await drain_logging_tasks()
assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE
assert usage.get_local(_OPENED_ON) == _OTHER_USAGE
assert (
await dual_cache.async_get_cache(key=_TURNED_ON)
usage.get_local(_TURNED_ON)
== _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN
)
@ -2057,14 +2058,12 @@ class TestWebSocketTurnAdmission:
) -> None:
import websockets.exceptions # noqa: F401 (lazy submodule must be importable)
from litellm.caching.dual_cache import DualCache
from litellm.responses.streaming_iterator import drain_logging_tasks
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
_turns_finish_at(monkeypatch, _OPENED_AT)
dual_cache: Final = DualCache()
opened_with: Final = await async_admit_io_token_turn(dual_cache, _ITPM_DEPLOYMENT, _OPENED_AT)
usage: Final = _memory_usage()
opened_with: Final = await async_admit_io_token_turn(usage, _ITPM_DEPLOYMENT, _OPENED_AT)
first, _, _ = _turn_frames()
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock()
@ -2073,7 +2072,7 @@ class TestWebSocketTurnAdmission:
websocket=_ClientSocket(),
backend_ws=backend_ws,
request_data={},
logging_obj=_connection_logged_through(_rate_limit_check(dual_cache), opened_with),
logging_obj=_connection_logged_through(_rate_limit_check(usage), opened_with),
first_message=first,
)
@ -2081,7 +2080,7 @@ class TestWebSocketTurnAdmission:
await handler.backend_to_client()
await drain_logging_tasks()
assert await dual_cache.async_get_cache(key=_ITPM_OPENED_ON) == counted_input_tokens
assert usage.get_local(_ITPM_OPENED_ON) == counted_input_tokens
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -2094,25 +2093,23 @@ class TestWebSocketTurnAdmission:
async def test_managed_relay_counts_every_turn_on_the_minute_it_finishes(
self, monkeypatch: pytest.MonkeyPatch, failing_turns: frozenset[int], first_turn_tokens: int
) -> None:
from litellm.caching.dual_cache import DualCache
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
_turns_finish_at(monkeypatch, _TURNED_AT)
dual_cache: Final = DualCache()
opened_with: Final = await _socket_opened_on(dual_cache)
usage: Final = _memory_usage()
opened_with: Final = await _socket_opened_on(usage)
await _run_managed_relay(
monkeypatch,
_TwoOutputTokenTurns(failing_turns),
_rate_limit_check(dual_cache),
_rate_limit_check(usage),
_turn_frames(),
_later_turns_admitted_through(dual_cache),
_later_turns_admitted_through(usage),
opened_with,
)
assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE
assert usage.get_local(_OPENED_ON) == _OTHER_USAGE
assert (
await dual_cache.async_get_cache(key=_TURNED_ON)
usage.get_local(_TURNED_ON)
== _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN
)
@ -2121,26 +2118,24 @@ class TestWebSocketTurnAdmission:
self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
_turns_finish_at(monkeypatch, _TURNED_AT)
provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=_provider_turn_stream)
client_metadata: Final = {"session": "abc", "team": "search"}
dual_cache: Final = DualCache()
usage: Final = _memory_usage()
await _relay_turns_to(
monkeypatch,
"openai/gpt-5-mini",
_rate_limit_check(dual_cache),
_rate_limit_check(usage),
tuple(
json.dumps({"type": "response.create", "input": word, "metadata": client_metadata})
for word in ("one", "two")
),
_later_turns_admitted_through(dual_cache),
await _socket_opened_on(dual_cache),
_later_turns_admitted_through(usage),
await _socket_opened_on(usage),
api_key="sk-fake",
)
@ -2148,7 +2143,7 @@ class TestWebSocketTurnAdmission:
client_metadata,
client_metadata,
]
assert await dual_cache.async_get_cache(key=_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN
assert usage.get_local(_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN
@pytest.mark.asyncio
async def test_native_relay_sends_the_provider_only_the_metadata_the_socket_opened_with(

15
tests/unit/test_router/core/test_registry_characterization_model_info.py
View file

@ -4,12 +4,15 @@ limits/mode/display name/service tiers/listing info, group usage this minute, re
the rate limit headers ``set_response_headers`` attaches.
"""
from typing import Final, Literal
import pytest
import litellm
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.router_core import RouterCore
from litellm.types.router import Deployment, RouterCacheEnum
from litellm.redis.stores.router_usage import UsageCounter
from litellm.types.router import Deployment
from litellm.utils import get_utc_datetime
from tests.unit.test_router.core.registry_helpers import MakeCore, typed, unique_id
from tests.unit.test_router.core.scripted_host import completion_response, deployment
@ -45,9 +48,9 @@ async def _serve(core: RouterCore, deployment_id: str, model: str, total_tokens:
)
async def _count(core: RouterCore, kind: RouterCacheEnum, deployment_id: str, model: str, value: int) -> None:
key = kind.value.format(id=deployment_id, model=model, current_minute=_minute())
await core.health().store.async_set_cache(key, value)
async def _count(core: RouterCore, kind: Literal["itpm", "otpm"], deployment_id: str, model: str, value: int) -> None:
usage: Final = core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the core counts io tokens in
usage.set_local(UsageCounter("global_router", deployment_id, model, _minute(), kind), value)
# -- model group info -----------------------------------------------------------------------------
@ -464,7 +467,7 @@ async def test_group_usage_sums_each_deployments_counters_for_this_minute(make_c
async def test_io_token_usage_sums_input_and_output_counters(make_core: MakeCore) -> None:
core = _usage_core(make_core)
assert await core.get_model_group_io_token_usage("io") == (None, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core
await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7)
await _count(core, "itpm", "i", "openai/gpt-4o", 7)
assert await core.get_model_group_io_token_usage("io") == (7, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core
@ -489,7 +492,7 @@ async def test_remaining_usage_of_an_io_limited_group_reports_input_and_output_t
make_core: MakeCore,
) -> None:
core = _usage_core(make_core)
await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7)
await _count(core, "itpm", "i", "openai/gpt-4o", 7)
remaining = await core.health().usage("io", kind="remaining")

3
tests/unit/test_router/test_enforce_model_rate_limits.py
View file

@ -14,7 +14,6 @@ import pytest
import litellm
from litellm import Router
from litellm.caching.dual_cache import DualCache
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_usage import CounterIncrement, CounterKind, RouterUsage, RouterUsageStore, UsageCounter
from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
@ -38,7 +37,7 @@ def _memory_usage() -> RouterUsageStore:
def _check(
usage: RouterUsage | None = None, clock: Callable[[], datetime] = get_utc_datetime
) -> ModelRateLimitingCheck:
return ModelRateLimitingCheck(usage=usage or _memory_usage(), dual_cache=DualCache(), clock=clock)
return ModelRateLimitingCheck(usage=usage or _memory_usage(), clock=clock)
def _counter(deployment: dict, kind: CounterKind, window: str) -> UsageCounter:

618
tests/unit/test_router/test_io_token_rate_limits.py
View file

File diff suppressed because it is too large Load diff

39
tests/unit/test_router/test_router_helper_utils.py
View file

@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.redis.stores.router_usage import UsageCounter
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
from litellm.types.litellm_params import RoutingStrategyName
from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
@ -347,8 +348,9 @@ def _io_limited_router(kind: str, **litellm_params: object) -> Router:
)
def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> str:
return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}"
def _minutes_count(router: Router, kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> object:
usage: Final = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in
return usage.get_local(UsageCounter("global_router", "ws-io-dep", model, minute.strftime("%H-%M"), kind))
_PROVIDER_RESPONSE: Final = {
@ -404,7 +406,7 @@ async def test_responses_call_sends_the_provider_only_the_client_metadata_and_se
assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata]
assert [
await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned))
_minutes_count(router, kind, router_minute_pinned)
for kind in ("itpm", "otpm")
] == [10, 5]
@ -509,7 +511,7 @@ async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attem
assert provider.call_count == 2
assert [
await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned, model))
_minutes_count(router, kind, router_minute_pinned, model)
for kind in ("itpm", "otpm")
] == [10, 5]
@ -925,7 +927,6 @@ async def test_get_model_group_io_token_usage_sums_across_deployments():
in the model group (not just the first), reading the same per-deployment
cache keys the pre-call reservation writes to.
"""
from litellm.types.router import RouterCacheEnum
from litellm.utils import get_utc_datetime
router = Router(
@ -952,26 +953,14 @@ async def test_get_model_group_io_token_usage_sums_across_deployments():
)
minute = get_utc_datetime().strftime("%H-%M")
keys_and_values = [
(
RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute),
30,
),
(
RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute),
10,
),
(
RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute),
70,
),
(
RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute),
20,
),
]
for key, value in keys_and_values:
await router.cache.async_increment_cache(key=key, value=value, ttl=60)
usage = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in
for deployment_id, model, kind, value in (
("io-usage-dep-1", "openai/gpt-4o-mini", "itpm", 30),
("io-usage-dep-1", "openai/gpt-4o-mini", "otpm", 10),
("io-usage-dep-2", "openai/gpt-4o", "itpm", 70),
("io-usage-dep-2", "openai/gpt-4o", "otpm", 20),
):
usage.set_local(UsageCounter("global_router", deployment_id, model, minute, kind), value)
current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus")