diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 6d5f7a65855..f6c86d75169 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -6,7 +6,7 @@ ###################################################################### import asyncio import os -from collections.abc import Mapping +from collections.abc import Mapping, MutableMapping from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast @@ -141,6 +141,20 @@ async def _raise_when_input_file_must_be_managed(model: str, credentials: Mappin ) +def _litellm_metadata_of(data: MutableMapping[str, object]) -> MutableMapping[str, object]: + """The request's litellm_metadata mapping, created on the request when it carries none. + + The success handler reads this mapping, so a flag or a model group set here has to live + inside it rather than beside it. + """ + existing: Final = data.get("litellm_metadata") + if isinstance(existing, MutableMapping): + return existing + created: Final[dict[str, object]] = {} # mutable-ok: the logging layer copies and extends this mapping + data["litellm_metadata"] = created # rebind-ok: the success handler reads the request's own mapping + return created + + def _raise_not_found_when_openai_fallback_unservable( requested_provider: "str | None", data: Mapping[str, object], @@ -668,11 +682,7 @@ async def retrieve_batch( poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active() if poller_owns_accounting: - litellm_metadata = data.get("litellm_metadata") - if not isinstance(litellm_metadata, dict): - litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend - data["litellm_metadata"] = litellm_metadata - litellm_metadata["batch_ignore_default_logging"] = True + _litellm_metadata_of(data)["batch_ignore_default_logging"] = True # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info @@ -697,6 +707,7 @@ async def retrieve_batch( # it the call falls into the legacy provider switch and 400s. data["model"] = model_from_id add_deployment_model_info(data=data, llm_router=llm_router, model_id=model_from_id) + _litellm_metadata_of(data).setdefault("model_group", model_from_id) # Retrieve batch using model credentials response = await litellm.aretrieve_batch( diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index bbfc7325f40..cfa54ae01a2 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -5,6 +5,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from openai.types import Batch + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -13,6 +15,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, StandardLoggingPayload @@ -117,6 +120,17 @@ def model_budget_start_time_cache_key( return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" +def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str: + return f"{spend_key}:batch:{batch_id}" + + +def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None: + """A finished batch reports its whole cost on every poll, so its id is charged once per counter.""" + if response_cost <= 0 or not is_batch_retrieve_call_type(call_type): + return None + return response_obj.id if isinstance(response_obj, Batch) else None + + def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None: """Find the `model_max_budget` entry that governs `model`, or None.""" for candidate in _budget_model_candidates(model): @@ -537,22 +551,18 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) return + batch_id: Final = batch_id_to_charge_once( + call_type=kwargs.get("call_type"), + response_obj=response_obj, + response_cost=response_cost, + ) for entity_type, entity_id, resolved in resolved_budgets: - await self._increment_spend_for_key( - budget_config=resolved.budget_config, - spend_key=model_budget_spend_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=resolved.budget_model, - budget_duration=resolved.budget_config.budget_duration, - ), - start_time_key=model_budget_start_time_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=resolved.budget_model, - budget_duration=resolved.budget_config.budget_duration, - ), + await self._charge_entity( + entity_type=entity_type, + entity_id=entity_id, + resolved=resolved, response_cost=response_cost, + batch_id=batch_id, ) if self.dual_cache.redis_cache is not None: @@ -562,3 +572,45 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): "current state of in memory cache %s", json.dumps(self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str), ) + + async def _charge_entity( + self, + entity_type: Litellm_EntityType, + entity_id: str | None, + resolved: ResolvedModelBudget, + response_cost: float, + batch_id: str | None, + ) -> None: + budget_duration: Final = resolved.budget_config.budget_duration + if budget_duration is None: + return + spend_key: Final = model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=budget_duration, + ) + if batch_id is not None and not await self._claim_batch_charge( + spend_key=spend_key, + batch_id=batch_id, + ttl_seconds=duration_in_seconds(budget_duration), + ): + return + await self._increment_spend_for_key( + budget_config=resolved.budget_config, + spend_key=spend_key, + start_time_key=model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=budget_duration, + ), + response_cost=response_cost, + ) + + async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: + marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) + polls: Final = await self.dual_cache.async_increment_cache( + key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True + ) + return polls == 1 diff --git a/litellm/router.py b/litellm/router.py index d5589a9b760..98c7c319eaa 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -152,6 +152,7 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, + is_batch_retrieve_call_type, replace_model_in_jsonl, should_replace_model_in_jsonl, ) @@ -6200,6 +6201,8 @@ class Router: """ try: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) + requested_model_group: Final = model + metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch") if model is not None: filtered_model_list: ( list[DeploymentTypedDict] | list[dict] | dict | None @@ -6236,6 +6239,9 @@ class Router: kwargs=new_kwargs, function_name="aretrieve_batch", ) + model_group: Final = requested_model_group or model_name["model_name"] + if not new_kwargs[metadata_variable_name].get("model_group"): + new_kwargs[metadata_variable_name]["model_group"] = model_group new_kwargs.pop("custom_llm_provider", None) data.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( @@ -7944,6 +7950,8 @@ class Router: # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object is None") @@ -8090,6 +8098,8 @@ class Router: - key: str - The key used to increment the cache - None: if no key is found """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return None id = None if kwargs["litellm_params"].get("metadata") is None: pass @@ -8218,6 +8228,8 @@ class Router: """ Update RPM usage for a deployment """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return deployment_name: Final = kwargs["litellm_params"]["metadata"].get( "deployment", None ) # handles wildcard routes - by giving the original name sent to `litellm.completion` diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 0b73f4e31a7..9ab670e4b95 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -9,6 +9,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 @@ -48,6 +49,8 @@ def _request_count_key(model_group: str, deployment_id: str) -> str: def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return None try: call: Final = _CALL_KWARGS.validate_python(kwargs) except ValidationError: diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index d271349914e..22c321c65fb 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -8,6 +8,7 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type class LowestCostLoggingHandler(CustomLogger): @@ -19,6 +20,8 @@ class LowestCostLoggingHandler(CustomLogger): self.router_cache = router_cache def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update usage on success @@ -92,6 +95,8 @@ class LowestCostLoggingHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update cost usage on success diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index e902192811c..66c8227195d 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -13,6 +13,7 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -58,6 +59,8 @@ class LowestLatencyLoggingHandler(CustomLogger): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update latency usage on success @@ -182,6 +185,8 @@ class LowestLatencyLoggingHandler(CustomLogger): """ Check if Timeout Error, if timeout set deployment latency -> 100 """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: metadata_field: Final = self._select_metadata_field(kwargs) _exception: Final = kwargs.get("exception", None) @@ -236,6 +241,8 @@ class LowestLatencyLoggingHandler(CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update latency usage on success diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 31c4b1d7e3f..d4abf1f8f70 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -8,6 +8,7 @@ from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose @@ -27,6 +28,8 @@ class LowestTPMLoggingHandler(CustomLogger): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success @@ -79,6 +82,8 @@ class LowestTPMLoggingHandler(CustomLogger): verbose_router_logger.debug(traceback.format_exc()) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 665ff69ab47..a2acce5fcb5 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -210,6 +211,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): return deployment # don't fail calls if eg. redis fails to connect def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM/RPM usage on success @@ -250,6 +253,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: """ Update TPM usage on success diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index ccb6ad95519..be20c358202 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -5,6 +5,7 @@ from typing import Final from litellm._logging import verbose_logger from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose +from litellm.types.utils import CallTypes class InMemoryFile(io.BytesIO): @@ -170,3 +171,21 @@ def _get_router_metadata_variable_name(function_name: str | None) -> str: return "litellm_metadata" else: return "metadata" + + +BATCH_RETRIEVE_CALL_TYPES: Final = frozenset( + { + CallTypes.aretrieve_batch.value, + CallTypes.retrieve_batch.value, + } +) + + +def is_batch_retrieve_call_type(call_type: object) -> bool: + """ + A batch retrieve reports the whole job's token usage, which the provider spent + asynchronously over the life of the batch, and reports it again on every poll of the + finished batch. The counters that measure live traffic, per-minute rate limits and the + routing strategies' own state, must not be fed from it. + """ + return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index c9f19731372..e274ac61a01 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -317,3 +317,18 @@ def test_replace_model_in_jsonl_with_embedded_newlines(): == "This is a message\nwith multiple\nlines" ) assert result_json["custom_id"] == "test123" + + +def test_is_batch_retrieve_call_type_matches_only_batch_retrieves(): + from litellm.router_utils.batch_utils import is_batch_retrieve_call_type + from litellm.types.utils import CallTypes + + assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True + assert is_batch_retrieve_call_type(CallTypes.retrieve_batch.value) is True + + for call_type in CallTypes: + if call_type in (CallTypes.aretrieve_batch, CallTypes.retrieve_batch): + continue + assert is_batch_retrieve_call_type(call_type.value) is False + + assert is_batch_retrieve_call_type(None) is False diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 8571ff20e57..2d597abf3b8 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1438,9 +1438,11 @@ async def call_retrieve( user: Optional[UserAPIKeyAuth] = None, headers: Optional[Dict[str, str]] = None, query: Optional[Dict[str, str]] = None, + enriched_data: Optional[Dict[str, Any]] = None, ): - # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...). - harness.data["data"] = {"batch_id": batch_id} + # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...), + # then pre-call enrichment adds key/team metadata to it. + harness.data["data"] = {"batch_id": batch_id, **(enriched_data or {})} return await endpoints.retrieve_batch( request=FakeRequest(headers=headers, query=query), fastapi_response=Response(), @@ -1476,6 +1478,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness): "api_key": "sk-azure", "api_base": "https://azure.test", "model": "azure/gpt-4o", + "litellm_metadata": {"model_group": "azure/gpt-4o"}, } # 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip. @@ -1522,6 +1525,39 @@ async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment assert retrieve_harness.aretrieve_kwargs()["model"] == "azure/gpt-4o" +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__stamps_model_group(retrieve_harness): + """This path never goes through the router, so nothing else labels the call. + Without the stamp the spend log lands under a blank model group and the batch + disappears from per-model usage.""" + await call_retrieve(retrieve_harness, AZURE_BATCH_ID) + + litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"] + + assert litellm_metadata["model_group"] == "azure/gpt-4o" + + +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__stamps_model_group_beside_existing_metadata( + retrieve_harness, +): + """The stamp joins the metadata pre-call enrichment already built. Replacing + that dict instead of adding to it drops the key and team labels the spend log + is attributed with.""" + await call_retrieve( + retrieve_harness, + AZURE_BATCH_ID, + enriched_data={"litellm_metadata": {"user_api_key_alias": "team-a-key"}}, + ) + + litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"] + + assert litellm_metadata == { + "user_api_key_alias": "team-a-key", + "model_group": "azure/gpt-4o", + } + + @pytest.mark.asyncio async def test_retrieve__model_encoded_id__encodes_output_and_error_ids( retrieve_harness, diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py new file mode 100644 index 00000000000..ffb60fb4651 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -0,0 +1,198 @@ +import asyncio +import time +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, +) +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.utils import LiteLLMBatch, Usage + +KEY_HASH: Final = "key-hash-batch" +USER_ID: Final = "user-batch" +MODEL_GROUP: Final = "batch-qa-primary" +BATCH_COST: Final = 2.925e-05 +CHAT_COST: Final = 0.001 +KEY_SPEND_KEY: Final = f"virtual_key_spend:{KEY_HASH}:{MODEL_GROUP}:1d" +USER_SPEND_KEY: Final = f"user_model_spend:{USER_ID}:{MODEL_GROUP}:1d" + + +def _batch(batch_id: str, status: str) -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-batch", + object="batch", + status=status, + usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38), + ) + + +def _event(call_type: str, response_cost: float) -> dict[str, object]: + return { + "call_type": call_type, + "standard_logging_object": { + "call_type": call_type, + "response_cost": response_cost, + "model": "openai/gpt-5.4-mini", + "model_group": MODEL_GROUP, + "metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_id": USER_ID}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}}, + "user_api_key_user_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}}, + } + }, + } + + +async def _poll(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBatch, response_cost: float) -> None: + await limiter.async_log_success_event( + _event("aretrieve_batch", response_cost), response_obj=batch, start_time=None, end_time=None + ) + + +async def _chat(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter) -> None: + await limiter.async_log_success_event( + _event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None + ) + + +async def _spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float: + return await limiter.dual_cache.async_get_cache(key=spend_key) or 0.0 + + +def _local_spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float: + return limiter.dual_cache.in_memory_cache.get_cache(key=spend_key) or 0.0 + + +class _Clock: + def __init__(self) -> None: + self.seconds = 0.0 + + def now(self) -> float: + return self.seconds + + def advance(self, seconds: float) -> None: + self.seconds = self.seconds + seconds + + +class _SharedRedisDouble: + def __init__(self, now: Callable[[], float] = time.time) -> None: + self.now = now + self.entries: Mapping[str, tuple[float, float | None]] = MappingProxyType({}) + + def _live(self, key: str) -> tuple[float, float | None] | None: + entry: Final = self.entries.get(key) + if entry is None: + return None + expires_at: Final = entry[1] + if expires_at is not None and expires_at <= self.now(): + return None + return entry + + def _store(self, key: str, value: float, expires_at: float | None) -> None: + self.entries = MappingProxyType({**self.entries, key: (value, expires_at)}) + + async def async_get_cache(self, key: str, **kwargs: object) -> float | None: + await asyncio.sleep(0) + entry: Final = self._live(key) + return None if entry is None else entry[0] + + async def async_set_cache(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> None: + await asyncio.sleep(0) + self._store(key, value, None if ttl is None else self.now() + ttl) + + async def async_increment( + self, + key: str, + value: float, + ttl: int | None = None, + parent_otel_span: object = None, + refresh_ttl: bool = False, + ) -> float: + await asyncio.sleep(0) + live: Final = self._live(key) + total: Final = value if live is None else live[0] + value + kept_expiry: Final = None if live is None else live[1] + expires_at: Final = ( + kept_expiry if ttl is None or (kept_expiry is not None and not refresh_ttl) else self.now() + ttl + ) + self._store(key, total, expires_at) + return total + + async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: + return [await self.async_increment(op["key"], op["increment_value"], ttl=op["ttl"]) for op in increment_list] + + +def _worker(redis: _SharedRedisDouble) -> _PROXY_VirtualKeyModelMaxBudgetLimiter: + return _PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(redis_cache=redis) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + + +async def _drain_redis_pushes() -> None: + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + + +@pytest.mark.asyncio +async def test_polls_of_a_finished_batch_charge_each_per_model_budget_once(): + limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + first: Final = _batch("batch_first", "completed") + + await _poll(limiter, _batch("batch_first", "in_progress"), response_cost=0) + for _ in range(3): + await _poll(limiter, first, response_cost=BATCH_COST) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + assert await _spend(limiter, USER_SPEND_KEY) == pytest.approx(BATCH_COST) + + +@pytest.mark.asyncio +async def test_a_second_batch_and_chat_requests_still_charge_the_budget(): + limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + + await _poll(limiter, _batch("batch_first", "completed"), response_cost=BATCH_COST) + await _poll(limiter, _batch("batch_first", "completed"), response_cost=BATCH_COST) + await _poll(limiter, _batch("batch_second", "completed"), response_cost=BATCH_COST) + await _chat(limiter) + await _chat(limiter) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(2 * BATCH_COST + 2 * CHAT_COST) + + +@pytest.mark.asyncio +async def test_two_workers_polling_the_same_finished_batch_at_once_charge_it_once(): + redis: Final = _SharedRedisDouble() + worker_a: Final = _worker(redis) + worker_b: Final = _worker(redis) + finished: Final = _batch("batch_first", "completed") + + await asyncio.gather(_poll(worker_a, finished, BATCH_COST), _poll(worker_b, finished, BATCH_COST)) + await _drain_redis_pushes() + + assert _local_spend(worker_a, KEY_SPEND_KEY) + _local_spend(worker_b, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + assert await redis.async_get_cache(KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + + +@pytest.mark.asyncio +async def test_a_batch_polled_within_every_budget_window_is_never_charged_again(): + clock: Final = _Clock() + limiter: Final = _worker(_SharedRedisDouble(now=clock.now)) + finished: Final = _batch("batch_first", "completed") + + await _poll(limiter, finished, BATCH_COST) + clock.advance(12 * 3600) + await _poll(limiter, finished, BATCH_COST) + clock.advance(18 * 3600) + await _poll(limiter, finished, BATCH_COST) + + assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 184092e4096..d62c9af7ee5 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -49,6 +49,7 @@ from litellm.router import ( from litellm.router_strategy import simple_shuffle from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.llms.openai import ChatCompletionRequest from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy @@ -1009,6 +1010,338 @@ async def test_arouter_aretrieve_batch(): assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base" +_BATCH_GROUP = "gemini-batch-group" +_BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini" +_BATCH_API_BASE = "http://localhost:4001/v1" +_BATCH_ID = "batch-1" +_BATCH_ROWS = 2 +_BATCH_TOKENS_PER_ROW = 600 + +_BATCH_COMPLETED = { + "id": _BATCH_ID, + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-1", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-1", + "error_file_id": None, + "created_at": 0, + "completed_at": 1, + "request_counts": {"total": _BATCH_ROWS, "completed": _BATCH_ROWS, "failed": 0}, + "metadata": None, +} + +_BATCH_OUTPUT_JSONL = "\n".join( + json.dumps( + { + "id": f"req-{row}", + "custom_id": f"row-{row}", + "response": { + "status_code": 200, + "body": { + "id": f"chatcmpl-{row}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": 500, + "completion_tokens": 100, + "total_tokens": _BATCH_TOKENS_PER_ROW, + }, + }, + }, + } + ) + for row in range(_BATCH_ROWS) +) + + +class _BatchPayloadCollector(CustomLogger): + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payloads.append(kwargs.get("standard_logging_object")) + + async def retrieve_batch_payload(self): + for _ in range(100): + for payload in self.payloads: + if payload and payload.get("call_type") == "aretrieve_batch": + return payload + await asyncio.sleep(0.05) + raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}") + + +def _batch_model_group_router(): + return litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + } + ] + ) + + +def _mock_batch_provider(respx_mock): + respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(200, json=_BATCH_COMPLETED) + ) + respx_mock.get(f"{_BATCH_API_BASE}/files/file-out-1/content").mock( + return_value=httpx.Response(200, text=_BATCH_OUTPUT_JSONL) + ) + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypatch: pytest.MonkeyPatch): + """ + The proxy retrieves a managed batch by id only - no `model` in the request. + The router fans out over its deployments, so the model group is only known + from the deployment that answered. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = _batch_model_group_router() + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + + assert response.id == _BATCH_ID + assert payload["total_tokens"] == _BATCH_ROWS * _BATCH_TOKENS_PER_ROW + assert payload["model"] == _BATCH_DEPLOYMENT_MODEL + assert payload["model_group"] == _BATCH_GROUP + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(monkeypatch: pytest.MonkeyPatch): + """An explicitly requested model group is what gets logged.""" + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = _batch_model_group_router() + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + await router.aretrieve_batch(model=_BATCH_GROUP, batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + + assert payload["model_group"] == _BATCH_GROUP + + +_UNRELATED_BATCH_GROUP = "unrelated-batch-group" +_UNRELATED_BATCH_API_BASE = "http://localhost:4002/v1" + +_BATCH_NOT_FOUND = { + "error": { + "message": f"No batch found with id '{_BATCH_ID}'.", + "type": "invalid_request_error", + "code": "batch_not_found", + } +} + + +async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:")) + if keys: + return keys + await asyncio.sleep(0.05) + return [] + + +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch): + """ + A batch reports the whole job's tokens on retrieve, and reports them again on every + poll of the finished batch, so they are not a measure of load in the current minute. + The fan-out also probes deployments the caller never named. Neither may reach the + per-minute tpm/rpm counters that gate live traffic. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + router = litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "batch-dep"}, + "tpm": 1000, + "rpm": 10, + }, + { + "model_name": _UNRELATED_BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _UNRELATED_BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "unrelated-dep"}, + "tpm": 1000, + "rpm": 10, + }, + ] + ) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + payload = await collector.retrieve_batch_payload() + usage_keys = await _router_usage_keys(router) + + assert response.id == _BATCH_ID + assert payload["model_group"] == _BATCH_GROUP + assert usage_keys == [] + + +@pytest.mark.parametrize( + ("call_type", "expected_key", "expected_successes"), + [ + ("aretrieve_batch", None, 0), + ("retrieve_batch", None, 0), + ("acompletion", "batch-dep:successes", 1), + ], +) +def test_sync_deployment_callback_on_success_skips_batch_retrieves( + call_type: str, expected_key: str | None, expected_successes: int +): + router = litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": {"model": _BATCH_DEPLOYMENT_MODEL, "api_base": _BATCH_API_BASE, "api_key": "sk-fake"}, + "model_info": {"id": "batch-dep"}, + } + ] + ) + + key = router.sync_deployment_callback_on_success( + kwargs={ + "call_type": call_type, + "litellm_params": {"metadata": {"model_group": _BATCH_GROUP}, "model_info": {"id": "batch-dep"}}, + }, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert key == expected_key + assert ( + get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep") + == expected_successes + ) + +_ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") + + +async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + cache_dict = router.cache.in_memory_cache.cache_dict + moved = sorted( + f"{key}={cache_dict[key]}" + for key in cache_dict + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) + and cache_dict[key] + ) + if moved: + return moved + await asyncio.sleep(0.05) + return [] + + +def _batch_fan_out_router(routing_strategy: str): + return litellm.Router( + routing_strategy=routing_strategy, + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "batch-dep"}, + }, + { + "model_name": _UNRELATED_BATCH_GROUP, + "litellm_params": { + "model": _BATCH_DEPLOYMENT_MODEL, + "api_base": _UNRELATED_BATCH_API_BASE, + "api_key": "sk-fake", + }, + "model_info": {"id": "unrelated-dep"}, + }, + ], + ) + + +@pytest.mark.parametrize( + "routing_strategy", + [ + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", + "least-busy", + ], +) +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( + monkeypatch: pytest.MonkeyPatch, routing_strategy: str +): + """ + Every routing strategy picks a deployment from what recent live traffic did. + A batch retrieve reports the whole job on every poll and probes deployments the + caller never named, so polling a finished batch must not move the numbers that + decide where the next chat request goes. + """ + import respx + + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + monkeypatch.setattr(litellm, "input_callback", []) + router = _batch_fan_out_router(routing_strategy) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + for _ in range(3): + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + await collector.retrieve_batch_payload() + moved_counters = await _moved_routing_counters(router) + + assert response.id == _BATCH_ID + assert moved_counters == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """