diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 05e3e0d839d..bd9440d4754 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -852,20 +852,26 @@ class CheckBatchCost: ) return None - try: - await logging_obj.async_success_handler( - result=response, - batch_cost=batch_result.cost, - batch_usage=batch_result.usage, - batch_models=batch_result.models, - batch_successful_requests=batch_result.successful_requests, - batch_failed_requests=batch_result.failed_requests, - batch_prompt_cost=batch_result.prompt_cost, - batch_completion_cost=batch_result.completion_cost, - ) - except Exception: - await self._release_job_claim(job, original_status, original_file_object) - raise + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + with track_batch_accounting(logging_obj.litellm_call_id) as accounting: + try: + await logging_obj.async_success_handler( + result=response, + batch_cost=batch_result.cost, + batch_usage=batch_result.usage, + batch_models=batch_result.models, + batch_successful_requests=batch_result.successful_requests, + batch_failed_requests=batch_result.failed_requests, + batch_prompt_cost=batch_result.prompt_cost, + batch_completion_cost=batch_result.completion_cost, + ) + finally: + if accounting.retryable_failure: + await self._release_job_claim(job, original_status, original_file_object) + + if accounting.retryable_failure: + return None # Record batch duration (completed_at - created_at) if prom_logger and response.completed_at and response.created_at: @@ -1186,7 +1192,7 @@ class CheckBatchCost: continue verbose_proxy_logger.error( f"CheckBatchCost: failed to track cost for batch {batch_id} " - f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}" + f"(job {job.id}): {tracking_err}" ) self._record_error(prom_logger, "cost_tracking_error") continue diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 00b5f9cd31c..37cef41034b 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -7,8 +7,6 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing -from openai.types import Batch - import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -19,7 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSet 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, budget_window_seconds -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, StandardLoggingPayload @@ -133,17 +131,6 @@ 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): diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 13a205d7e9f..1f8b352fed9 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,6 +1,8 @@ import asyncio import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, cast @@ -68,6 +70,37 @@ from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging + +class BatchAccountingOutcome: + def __init__(self, call_id: str) -> None: + self.call_id: Final = call_id + self.write_started = False + self.failed = False + + @property + def retryable_failure(self) -> bool: + return self.failed and not self.write_started + + +_BATCH_ACCOUNTING: Final[ContextVar[BatchAccountingOutcome | None]] = ContextVar("batch_accounting", default=None) +_BATCH_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) + + +@contextmanager +def track_batch_accounting(call_id: str) -> Generator[BatchAccountingOutcome, None, None]: + outcome: Final = BatchAccountingOutcome(call_id) + token: Final = _BATCH_ACCOUNTING.set(outcome) + try: + yield outcome + finally: + _BATCH_ACCOUNTING.reset(token) + + +def _batch_accounting_for(call_id: object) -> BatchAccountingOutcome | None: + outcome: Final = _BATCH_ACCOUNTING.get() + return outcome if outcome is not None and call_id == outcome.call_id else None + + _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { CallTypes.pass_through.value, @@ -445,6 +478,11 @@ class _ProxyDBLogger(CustomLogger): update_cache, ) + accounting: Final = ( + _batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id")) + if _BATCH_ACCOUNTING.get() is not None + else None + ) verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") try: verbose_proxy_logger.debug( @@ -465,6 +503,8 @@ class _ProxyDBLogger(CustomLogger): ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + if accounting is not None and (budget_reservation is not None or litellm.apply_search_free_tier): + accounting.write_started = True if ( isinstance(completion_response, LiteLLMBatch) and kwargs.get("call_type") == CallTypes.aretrieve_batch.value @@ -617,6 +657,8 @@ class _ProxyDBLogger(CustomLogger): f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" ) except Exception as e: + if accounting is not None: + accounting.failed = True failing_model: Final = kwargs.get("model", "") failing_call_type: Final = kwargs.get("call_type", "") error_msg: Final = ( @@ -936,7 +978,15 @@ async def _update_database_and_spend_counters_in_batch( project_id: str | None, ) -> bool: try: - charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( + writer: Final = proxy_logging_obj.db_spend_update_writer + accounting: Final = ( + _batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id")) + if _BATCH_ACCOUNTING.get() is not None + else None + ) + if accounting is not None: + accounting.write_started = True + charged: Final = await writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 7573bb6716f..147ce60a99b 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -28,6 +28,8 @@ from itertools import groupby from types import MappingProxyType from typing import Any, Final +from pydantic import TypeAdapter + import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger @@ -39,6 +41,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs +from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, ) @@ -48,6 +51,8 @@ from litellm.types.utils import BudgetConfig, GenericBudgetConfigType, StandardL from litellm.types.utils import BudgetConfig as GenericBudgetInfo DEFAULT_REDIS_SYNC_INTERVAL: Final = 1 +_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) +_BATCH_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) class _LiteLLMParamsDictView: @@ -505,16 +510,21 @@ class RouterBudgetLimiting(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") - if kwargs.get("call_type") == "_arealtime": + callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs) + if callback_fields.get("call_type") == "_arealtime": return standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: raise ValueError("standard_logging_payload is required") - response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + batch_response, response_cost = _BATCH_CHARGE_VALUES.validate_python( + (response_obj, _CALLBACK_FIELDS.validate_python(standard_logging_payload).get("response_cost", 0)) + ) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") + batch_id: Final = batch_id_to_charge_once(callback_fields.get("call_type"), batch_response, response_cost) + budget_config: Final = ( self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None ) @@ -527,6 +537,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=spend_key, start_time_key=start_time_key, response_cost=response_cost, + batch_id=batch_id, ) deployment_budget_config: Final = self._get_budget_config_for_deployment(model_id) @@ -539,11 +550,12 @@ class RouterBudgetLimiting(CustomLogger): spend_key=deployment_spend_key, start_time_key=deployment_start_time_key, response_cost=response_cost, + batch_id=batch_id, ) request_tags: Final = get_tags_from_request_kwargs( - kwargs, - metadata_variable_name=get_metadata_variable_name_from_kwargs(kwargs or {}), + callback_fields, + metadata_variable_name=get_metadata_variable_name_from_kwargs(callback_fields), ) if len(request_tags) > 0: for _tag in request_tags: @@ -556,6 +568,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=_tag_spend_key, start_time_key=_tag_start_time_key, response_cost=response_cost, + batch_id=batch_id, ) async def _increment_spend_for_key( @@ -564,12 +577,15 @@ class RouterBudgetLimiting(CustomLogger): spend_key: str, start_time_key: str, response_cost: float, + batch_id: str | None = None, ): if budget_config.budget_duration is None: return current_time: Final = self._clock() new_window_seconds: Final = budget_window_seconds(current_time, budget_config.budget_duration) + if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, new_window_seconds): + return budget_start: Final = await self._get_or_set_budget_start_time( start_time_key=start_time_key, @@ -595,6 +611,17 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, 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) + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is not None: + polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True) + return polls == 1 + local_polls: Final = await self.dual_cache.in_memory_cache.async_increment( + key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True + ) + return local_polls == 1 + @with_service_target("router_budgets") async def periodic_sync_in_memory_spend_with_redis(self): """ diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 88c58412e01..5b94d4b4e5b 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -6,6 +6,7 @@ from contextvars import ContextVar from os import PathLike from typing import Final, TypeVar +from openai.types import Batch from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger @@ -286,3 +287,14 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: routing strategies' own state, must not be fed from it. """ return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES + + +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 diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 241cf90eb5d..33b8085e082 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3104,6 +3104,151 @@ class TestMultiPodBatchCostClaim: ): yield logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external")) + @pytest.mark.parametrize("collector_enabled", (False, True)) + async def test_builtin_accounting_retries_only_before_a_possible_charge( + self, failure_phase: str, collector_enabled: bool + ) -> None: + import base64 + from datetime import datetime + + import httpx + import litellm + import respx + + from litellm.caching.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.utils import PrismaClient + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + from litellm.types.utils import SpecialEnums + + class AccountingDependency: + def __init__(self) -> None: + self.failure_phase = failure_phase + self.writer: Final = DBSpendUpdateWriter() + self.failed_tracking_alert = AsyncMock() + self.slack_alerting_instance = MagicMock() + self.slack_alerting_instance.customer_spend_alert = AsyncMock( + side_effect=ConnectionError("Alert destination unavailable after charge") + if failure_phase == "after_write" + else None + ) + + @property + def db_spend_update_writer(self) -> DBSpendUpdateWriter: + if self.failure_phase == "before_write": + raise ConnectionError("Accounting dependency temporarily unavailable") + return self.writer + + class ExternalCallback(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if failure_phase == "external": + raise ConnectionError("External destination unavailable") + + budget_cache: Final = DualCache() + RouterBudgetLimiting( + dual_cache=budget_cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "batch-alias", + "litellm_params": {"model": "openai/batch-test-model", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "model-123"}, + } + ], + ) + dependency: Final = AccountingDependency() + spend_prisma: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost:5432/unused", proxy_logging_obj=dependency + ) + create_spend_rows: Final = AsyncMock(return_value=1) + input_rate: Final = 0.000002 + litellm.register_model( + { + "batch-test-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": input_rate, + "output_cost_per_token_batches": 0.000006, + } + } + ) + content: Final = json.dumps( + { + "response": { + "status_code": 200, + "body": { + "model": "batch-test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 0, "total_tokens": 100}, + }, + }, + "error": None, + } + ).encode() + row: Final = _FakeManagedObjectRow() + original_status: Final = row.status + original_file_object: Final = row.file_object + row.unified_object_id = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode() + ) + .decode() + .rstrip("=") + ) + prisma: Final = self._prisma(row, []) + router: Final = self._router() + router.get_deployment.return_value = Deployment( + model_name="batch-alias", + litellm_params=LiteLLM_Params(model="openai/batch-test-model"), + model_info=ModelInfo(id="model-123"), + ) + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer + + producer: Final = MagicMock(spec=SpendEventProducer) if collector_enabled else None + litellm.logging_callback_manager.add_litellm_async_success_callback( + ProxyDBLogger(spend_event_producer=producer) + ) + litellm.logging_callback_manager.add_litellm_async_success_callback(ExternalCallback()) + with ( + respx.mock, + patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency), + patch("litellm.proxy.proxy_server.prisma_client", spend_prisma), + patch.object( + spend_prisma.writer_db._original_prisma, + "litellm_spendlogs", + MagicMock(create_many=create_spend_rows), + ), + ): + respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock( + return_value=httpx.Response(200, content=content) + ) + worker: Final = self._instance(prisma, router) + await worker.check_batch_cost() + assert row.batch_processed is (failure_phase != "before_write") + if failure_phase == "before_write": + assert row.status == original_status + assert json.loads(row.file_object) == original_file_object + assert create_spend_rows.await_count == (0 if failure_phase == "before_write" else 1) + dependency.failure_phase = "recovered" + await worker.check_batch_cost() + await worker.check_batch_cost() + + assert row.batch_processed is True + assert await budget_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(100 * input_rate) + assert await budget_cache.async_get_cache("deployment_spend:model-123:1d") == pytest.approx(100 * input_rate) + create_spend_rows.assert_awaited_once() + assert create_spend_rows.call_args.kwargs["data"][0]["spend"] == pytest.approx(100 * input_rate) + if producer is not None: + assert all( + json.loads(call.args[0])["call_type"] != "aretrieve_batch" for call in producer.publish.call_args_list + ) + @pytest.mark.asyncio @pytest.mark.parametrize( ("deployment_model", "successes", "failures", "expected_model"), @@ -3389,30 +3534,23 @@ class TestMultiPodBatchCostClaim: @pytest.mark.asyncio @pytest.mark.parametrize("stored_as_json_string", (True, False)) - async def test_a_failed_spend_log_write_releases_the_claim(self, stored_as_json_string: bool) -> None: - """A transient failure while billing a claimed batch must hand the row back, or its - spend is silently lost instead of being retried on the next cycle.""" - row = _FakeManagedObjectRow() - original_file_object: Final = row.file_object if stored_as_json_string else json.loads(row.file_object) - row.file_object = original_file_object - journal = [] - prisma = self._prisma(row, journal) + async def test_an_unclassified_logging_failure_keeps_the_claim(self, stored_as_json_string: bool) -> None: + row: Final = _FakeManagedObjectRow() + row.file_object = row.file_object if stored_as_json_string else json.loads(row.file_object) + journal: Final = [] + prisma: Final = self._prisma(row, journal) - with self._billing_patches(journal, bill_error=Exception("spend log write failed")): + with self._billing_patches(journal, bill_error=Exception("Accounting outcome unknown")): await self._instance(prisma, self._router()).check_batch_cost() - assert journal == ["fetch", "claim", "bill", "release"] - assert row.batch_processed is False - assert row.status == "in_progress" - assert json.loads(row.file_object) == original_file_object + assert journal == ["fetch", "claim", "bill"] + assert row.batch_processed is True with self._billing_patches(journal) as logging_obj: await self._instance(prisma, self._router()).check_batch_cost() await self._instance(prisma, self._router()).check_batch_cost() - assert row.batch_processed is True - assert row.status == "complete" - logging_obj.async_success_handler.assert_awaited_once() + logging_obj.async_success_handler.assert_not_awaited() @pytest.mark.asyncio async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self): diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 6be4b93c173..289f8cd4a1a 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -3270,3 +3270,100 @@ async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when else: assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown assert counter.totals == {} + + +@pytest.mark.asyncio +async def test_batch_accounting_failures_stay_with_their_originating_call() -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + logger: Final = ProxyDBLogger() + dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock()) + + async def fail_before_accounting(call_id: str) -> None: + await logger.async_log_success_event( + kwargs={ + "litellm_call_id": call_id, + "call_type": "aretrieve_batch", + "model": "batch-model", + "stream": False, + "response_cost": None, + }, + response_obj=None, + start_time=_CALL_START, + end_time=_CALL_START, + ) + + async def independent_batch() -> bool: + with track_batch_accounting("independent") as outcome: + await fail_before_accounting("independent") + return outcome.retryable_failure + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency): + with track_batch_accounting("outer") as outer: + await fail_before_accounting("unrelated") + assert outer.retryable_failure is False + with track_batch_accounting("inner") as inner: + await fail_before_accounting("inner") + assert inner.retryable_failure is True + assert outer.retryable_failure is False + independent, _ = await asyncio.gather(independent_batch(), fail_before_accounting("outer")) + assert independent is True + assert outer.retryable_failure is True + with pytest.raises(asyncio.CancelledError): + with track_batch_accounting("cancelled") as cancelled: + raise asyncio.CancelledError + await fail_before_accounting("cancelled") + assert cancelled.retryable_failure is False + + +@pytest.mark.asyncio +async def test_batch_accounting_does_not_retry_after_settling_a_budget_reservation() -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": ()} + dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock()) + with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency): + with track_batch_accounting("reserved") as outcome: + await ProxyDBLogger().async_log_success_event( + kwargs={ + "litellm_call_id": "reserved", + "call_type": "aretrieve_batch", + "model": "batch-model", + "stream": False, + "response_cost": None, + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + }, + response_obj=None, + start_time=_CALL_START, + end_time=_CALL_START, + ) + assert reservation["finalized"] is True + assert outcome.retryable_failure is False + + +@pytest.mark.asyncio +async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + monkeypatch.setattr(litellm, "apply_search_free_tier", True) + counter: Final = _SharedCounter() + dependency: Final = SimpleNamespace( + failed_tracking_alert=AsyncMock(), + internal_usage_cache=SimpleNamespace(dual_cache=SimpleNamespace(redis_cache=counter)), + ) + kwargs: Final = {**_grounded_gemini_kwargs(response_cost=1.0), "litellm_call_id": "grounded-batch"} + response: Final = ModelResponse( + usage=Usage( + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=5), + ) + ) + with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency): + with track_batch_accounting("grounded-batch") as outcome: + await ProxyDBLogger().async_log_success_event(kwargs, response, _CALL_START, _CALL_START) + + assert sum(counter.totals.values()) == 5 + assert outcome.failed is True + assert outcome.retryable_failure is False diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 3c825bd560f..c0583566531 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -449,3 +449,91 @@ def cleanup_redis(): redis_client.delete(key) except Exception as e: print(f"Error cleaning up Redis: {str(e)}") + + +def _completed_budget_batch(batch_id: str): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-budget", + object="batch", + status="completed", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ("aretrieve_batch", "retrieve_batch")) +async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type): + monkeypatch.setattr(litellm, "tag_budget_config", {"release": {"budget_limit": 100, "time_period": "1d"}}) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "batch-budget", + "litellm_params": {"model": "openai/batch-budget", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + kwargs: Final = { + **_success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type), + "metadata": {"tags": ["release"]}, + } + response: Final = _completed_budget_batch("batch-first") + await asyncio.gather(*(limiter.async_log_success_event(kwargs, response, None, None) for _ in range(10))) + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d", "tag_spend:release:1d") + assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.25, 0.25, 0.25) + await limiter.async_log_success_event(kwargs, _completed_budget_batch("batch-second"), None, None) + await limiter.async_log_success_event({**kwargs, "call_type": "acompletion"}, response, None, None) + await limiter.async_log_success_event({**kwargs, "call_type": "acompletion"}, response, None, None) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) + await limiter.async_log_success_event({**kwargs, "call_type": "_arealtime"}, response, None, None) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) + + +@pytest.mark.asyncio +async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): + from fakeredis import FakeServer, FakeRedis as FakeSyncRedis + from redis.exceptions import ConnectionError as RedisConnectionError + from fakeredis.aioredis import FakeRedis + from litellm.caching.redis_cache import RedisCache + + server: Final = FakeServer() + client: Final = FakeRedis(server=server) + sync_client: Final = FakeSyncRedis(server=server) + redis_cache: Final = RedisCache(host="router-batch-fixture", connection_pool=sync_client.connection_pool) + litellm.in_memory_llm_clients_cache.set_cache(key=redis_cache._get_async_client_cache_key(), value=client) + first: Final = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + ) + second: Final = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + ) + kwargs: Final = _success_kwargs( + provider_in_litellm_params=None, provider_in_payload="openai", call_type="aretrieve_batch" + ) + response: Final = _completed_budget_batch("batch-shared") + await asyncio.gather( + first.async_log_success_event(kwargs, response, None, None), + second.async_log_success_event(kwargs, response, None, None), + ) + await first._push_in_memory_increments_to_redis() + await second._push_in_memory_increments_to_redis() + assert float(await client.get("provider_spend:openai:1d")) == 0.25 + server.connected = False + with pytest.raises(RedisConnectionError): + await second.async_log_success_event(kwargs, _completed_budget_batch("batch-unconfirmed"), None, None) + server.connected = True + await second.async_log_success_event(kwargs, _completed_budget_batch("batch-unconfirmed"), None, None) + await second._push_in_memory_increments_to_redis() + assert float(await client.get("provider_spend:openai:1d")) == 0.5 + await client.aclose() diff --git a/tests/unit/router_utils/test_batch_utils.py b/tests/unit/router_utils/test_batch_utils.py index 4aa503ee528..79fe86a37c0 100644 --- a/tests/unit/router_utils/test_batch_utils.py +++ b/tests/unit/router_utils/test_batch_utils.py @@ -2,7 +2,7 @@ import asyncio import importlib import json from io import BytesIO -from typing import Dict, List +from typing import Dict, Final, List import pytest @@ -380,3 +380,31 @@ def setup_and_teardown(): yield loop.close() asyncio.set_event_loop(None) + + +@pytest.mark.parametrize( + ("call_type", "cost", "eligible"), + (("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False), + ("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)), +) +def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: + from openai.types import Batch + from litellm.router_utils.batch_utils import batch_id_to_charge_once + + batch: Final = Batch( + id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", + input_file_id="file-budget", object="batch", status="completed", + ) + assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None) + assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None + + +def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.router_utils.batch_utils import batch_charged_once_marker_key + + cache: Final = InMemoryCache() + cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) + assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1 + assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None + assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None