diff --git a/litellm/router.py b/litellm/router.py index 2397c6b2fc7..8f913d96463 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6177,7 +6177,8 @@ class Router: function_name="aretrieve_batch", ) model_group: Final = requested_model_group or model_name["model_name"] - new_kwargs[metadata_variable_name].setdefault("model_group", model_group) + 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( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..e93288fd9fe 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -11,6 +11,7 @@ from typing import Final 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 LeastBusyLoggingHandler(CustomLogger): @@ -27,6 +28,8 @@ class LeastBusyLoggingHandler(CustomLogger): Caching based on model group. """ + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -48,6 +51,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass def log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -76,6 +81,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass def log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass @@ -103,6 +110,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass 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: if kwargs["litellm_params"].get("metadata") is None: pass @@ -131,6 +140,8 @@ class LeastBusyLoggingHandler(CustomLogger): pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return try: if kwargs["litellm_params"].get("metadata") is None: pass diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index b927df0c438..aaad6186484 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 @@ -96,6 +99,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 a1b67eaeaf9..598ca1227ec 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -9,6 +9,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: @@ -35,6 +36,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 @@ -167,6 +170,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) @@ -221,6 +226,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_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 6e110b586fb..be20c358202 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -185,6 +185,7 @@ 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. Per-minute usage counters must not be fed from it. + 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/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d8045722998..836085c1f5d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1166,6 +1166,85 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m assert usage_keys == [] +_ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") + + +async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + keys = sorted( + key + for key in router.cache.in_memory_cache.cache_dict + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) + ) + if keys: + return keys + 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", "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() + strategy_keys = await _router_strategy_keys(router) + + assert response.id == _BATCH_ID + assert strategy_keys == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """