diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a3227bca235..14b9c48b0d6 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -500,6 +500,7 @@ class RequestRateLimiterStash: reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) + fallback_models: frozenset[str] = field(default_factory=frozenset) ptu_ceiling: PTUTeamCeiling | None = None ptu_reserved_tokens: int = 0 ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset) @@ -3786,6 +3787,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) -> bool: """Reserve the caller's per-model limits on a router fallback target. False when any is exhausted.""" data: Final = MappingProxyType({**request_data, "model": model}) + stash: Final = _request_stash.get() hidden_stash: Final = _request_stash.set(None) try: built: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None) @@ -3794,10 +3796,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors: Final = [ descriptor for descriptor in built if descriptor["key"] in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS ] - if not descriptors: - return True - response: Final = await self.should_rate_limit(descriptors, parent_otel_span=user_api_key_dict.parent_otel_span) - return response["overall_code"] == "OK" + allowed: Final = ( + not descriptors + or (await self.should_rate_limit(descriptors, parent_otel_span=user_api_key_dict.parent_otel_span))[ + "overall_code" + ] + == "OK" + ) + if allowed and stash is not None: + stash.fallback_models = stash.fallback_models | {model} + return allowed @asynccontextmanager async def request_capacity( @@ -4720,12 +4728,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens: Final = open_reservation.reserved_tokens if open_reservation is not None else 0 reserved_model: Final = stash.reserved_model if stash is not None else None reserved_scopes: Final[frozenset[tuple[str, str]]] = stash.reserved_scopes if stash is not None else frozenset() - # Reconciliation must target the same model-scoped counter that the - # pre-call reservation incremented. If a reservation was made, + served_model: Final = self._rate_limited_model(model_group) # ``reserved_model`` (resolved at admission, so an alias map reload - # mid-flight cannot move the charge) is authoritative; otherwise fall - # back to the router's ``model_group`` (the no-reservation charge path). - reconcile_model: Final = reserved_model if reserved_model is not None else self._rate_limited_model(model_group) + # mid-flight cannot move the charge) settles the request unless the + # router fell back to a group this limiter admitted as a fallback + # target: then the served group is charged and the abandoned group's + # per-model reservation is refunded. Client-supplied ``model_group`` + # metadata alone never moves the charge. + fell_back: Final = ( + stash is not None + and reserved_model is not None + and served_model is not None + and served_model.requested != reserved_model.requested + and served_model.requested in stash.fallback_models + ) + reconcile_model: Final = served_model if fell_back or reserved_model is None else reserved_model reconcile_models: Final = self._rate_limited_models(reconcile_model) pipeline_operations: Final[list[ReservationOperation]] = [] @@ -4764,6 +4781,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) charged_targets: Final = [target for target in targets if target not in key_owned_team_scopes] + abandoned_scopes: Final = ( + tuple( + sorted( + scope + for scope in reserved_scopes + if scope[0] in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS and scope not in charged_targets + ) + ) + if fell_back + else () + ) if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( "Releasing unused TPM budget on success: reserved=%s, actual=%s, release=%s", @@ -4779,6 +4807,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=reserved_tokens, ) ) + pipeline_operations.extend( + self._build_reservation_aware_tpm_ops( + targets=abandoned_scopes, + reserved_scopes=reserved_scopes, + actual_tokens=0, + reserved_tokens=reserved_tokens, + ) + ) call_type: Final[object] = kwargs.get("call_type") ptu_skips_socket_usage: Final = call_type == CallTypes.aresponses_websocket.value pipeline_operations.extend( diff --git a/tests/integration/routing/test_fallback_settles_serving_model.py b/tests/integration/routing/test_fallback_settles_serving_model.py new file mode 100644 index 00000000000..b9c32f7e91f --- /dev/null +++ b/tests/integration/routing/test_fallback_settles_serving_model.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import os +import time +import uuid +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedState, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import wire_server +from integration.routing.test_fallback_rate_limit_check import provider_error_reply +from integration.routing.test_per_model_rate_limit_fallbacks import ( + assert_served, + chat, + healthy_reply, + model_entry, + proxy_config, +) +from redis import Redis + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +ADMISSION_PROBE_TOKENS: Final = 1 +SETTLEMENT_DEADLINE_SECONDS: Final = 15.0 + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + primary: str + fallback: str + + +def key_model_tokens(group: str) -> int: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2) as cache: + counters: Final = tuple(cache.scan_iter(match=f"*model_per_key:*:{group}}}:tokens")) + return sum(int(cache.get(counter) or 0) for counter in counters) + + +def settled_key_model_tokens(primary: str, fallback: str, total_tokens: int) -> tuple[int, int]: + deadline: Final = time.monotonic() + SETTLEMENT_DEADLINE_SECONDS + while True: + charged: Final = (key_model_tokens(primary), key_model_tokens(fallback)) + if (charged[0] == 0 and charged[1] >= total_tokens) or time.monotonic() >= deadline: + return charged + time.sleep(0.5) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:10] + primary: Final = f"settle-primary-{suffix}" + fallback: Final = f"settle-fallback-{suffix}" + directory: Final = tmp_path_factory.mktemp("fallback-settles-serving-model") + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + primary_wire: Final = stack.enter_context(wire_server(provider_error_reply)) + fallback_wire: Final = stack.enter_context(wire_server(healthy_reply)) + config: Final = proxy_config( + directory, + model_list=(model_entry(primary, primary_wire), model_entry(fallback, fallback_wire)), + fallbacks=({primary: [fallback]},), + general_settings={}, + ) + owned: Final = stack.enter_context( + owned_proxy_process( + rig_gateway, directory, {}, config=config, workers=2, state=OwnedState.SHARED_WITH_WORKER + ) + ) + yield Rig(owned.gateway, primary, fallback) + + +def test_provider_error_fallback_charges_its_tokens_to_the_model_that_served(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_tpm_limit={rig.primary: 100_000, rig.fallback: 1_000}) + + response: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}") + assert_served(response, rig.fallback, "served pmrl-") + total_tokens: Final = int(response.json()["usage"]["total_tokens"]) + assert total_tokens > 0, response.text + + primary_tokens, fallback_tokens = settled_key_model_tokens(rig.primary, rig.fallback, total_tokens) + assert primary_tokens == 0, (primary_tokens, fallback_tokens) + assert total_tokens <= fallback_tokens <= total_tokens + ADMISSION_PROBE_TOKENS, ( + primary_tokens, + fallback_tokens, + ) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index a39d02a2788..e3a63d097ec 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -9941,3 +9941,144 @@ async def test_end_user_model_rpm_limit_is_shared_across_keys_and_isolated_per_c await _rpm_request(handler, cache, globex, "target") await _rpm_request(handler, cache, acme_via_key_one, "other-model") + + + +@pytest.mark.asyncio +async def test_success_tpm_accounting_moves_the_charge_to_the_model_group_that_served_a_fallback() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key_metadata: Final = {"model_tpm_limit": {"primary": 1000, "fallback": 1000}} + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-fallback-settle"), metadata=key_metadata) + + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "primary", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10}, + call_type="acompletion", + ) + stash: Final = get_request_stash() + assert stash is not None + assert stash.reserved_model == RateLimitedModel(requested="primary", group="primary") + assert stash.reserved_tokens > 0 + assert await handler.async_fallback_target_allowed( + auth, "fallback", {"messages": [{"role": "user", "content": "hello"}]} + ) + + response: Final = ModelResponse( + id="fallback-settle", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="fallback", + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + choices=[], + ) + kwargs: Final = { + "standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}}, + "litellm_params": { + "metadata": { + "model_group": "fallback", + "original_model_group": "primary", + "user_api_key_metadata": key_metadata, + } + }, + "model": "fallback", + } + + ops: Final = handler._build_success_event_pipeline_operations( + kwargs=kwargs, response_obj=response, rate_limit_type="total" + ) + + charged: Final = {_increment_of(op).counter: _increment_of(op).amount for op in ops} + assert charged[_counter("model_per_key", f"{auth.api_key}:fallback")] == 150 + assert charged[_counter("model_per_key", f"{auth.api_key}:primary")] == -stash.reserved_tokens + + +@pytest.mark.asyncio +async def test_success_tpm_accounting_settles_an_alias_to_target_fallback_once_per_counter( + collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str], +) -> None: + cache, handler, alias, target = collision_limiter + key_metadata: Final = {"model_tpm_limit": {alias: 1000, target: 1000}} + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-alias-fallback-settle"), metadata=key_metadata) + + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": alias, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10}, + call_type="acompletion", + ) + stash: Final = get_request_stash() + assert stash is not None + assert stash.reserved_model == RateLimitedModel(requested=alias, group=target) + assert stash.reserved_tokens > 0 + assert await handler.async_fallback_target_allowed( + auth, target, {"messages": [{"role": "user", "content": "hello"}]} + ) + + response: Final = ModelResponse( + id="alias-fallback-settle", + object="chat.completion", + created=int(datetime.now().timestamp()), + model=target, + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + choices=[], + ) + kwargs: Final = { + "standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}}, + "litellm_params": { + "metadata": {"model_group": target, "original_model_group": alias, "user_api_key_metadata": key_metadata} + }, + "model": target, + } + + ops: Final = handler._build_success_event_pipeline_operations( + kwargs=kwargs, response_obj=response, rate_limit_type="total" + ) + + increments: Final = tuple(_increment_of(op) for op in ops) + target_counter: Final = _counter("model_per_key", f"{auth.api_key}:{target}") + alias_counter: Final = _counter("model_per_key", f"{auth.api_key}:{alias}") + assert [inc.amount for inc in increments if inc.counter == target_counter] == [150 - stash.reserved_tokens] + assert [inc.amount for inc in increments if inc.counter == alias_counter] == [] + + + +@pytest.mark.asyncio +async def test_success_tpm_accounting_ignores_a_client_supplied_model_group_without_a_router_fallback() -> None: + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key_metadata: Final = {"model_tpm_limit": {"primary": 1000}} + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-client-model-group"), metadata=key_metadata) + + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "primary", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10}, + call_type="acompletion", + ) + stash: Final = get_request_stash() + assert stash is not None + assert stash.reserved_tokens > 0 + + response: Final = ModelResponse( + id="client-model-group", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="primary", + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + choices=[], + ) + kwargs: Final = { + "standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}}, + "litellm_params": {"metadata": {"model_group": "uncapped", "user_api_key_metadata": key_metadata}}, + "model": "primary", + } + + ops: Final = handler._build_success_event_pipeline_operations( + kwargs=kwargs, response_obj=response, rate_limit_type="total" + ) + + charged: Final = {_increment_of(op).counter: _increment_of(op).amount for op in ops} + assert charged[_counter("model_per_key", f"{auth.api_key}:primary")] == 150 - stash.reserved_tokens + assert not any(":uncapped" in counter.scope.descriptor_value for counter in charged)