From f585b558e95fae00b18cf7cd0058a508b98cecc7 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 19 Oct 2026 18:51:51 -0700 Subject: [PATCH] fix(batches): apply the Vertex regional endpoint uplift to batch cost (#47466) * fix(batches): apply the Vertex regional endpoint uplift to batch cost batch_cost_calculator takes vertex_location and multiplies by the model's regional_endpoint_uplift_multiplier through get_vertex_regional_endpoint_uplift, and every batch cost path (aretrieve_batch pricing, completed-batch logging, native Vertex batch rows, the proxy batch cost poll and its per-line items) resolves the deployment's location and passes it down Co-authored-by: Mrun25 <165297479+Mrun25@users.noreply.github.com> * test(logging): point the resolver tests at resolve_vertex_location_for_cost --------- Backport of #47466 to stable/1.104.x Cherry-picked from c6c303a6181bf0b59f6b30beb708deade0d955d0 (main) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Mrun25 <165297479+Mrun25@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 12 ++++- litellm/batches/batch_utils.py | 34 ++++++++++-- litellm/cost_calculator.py | 7 ++- litellm/litellm_core_utils/litellm_logging.py | 10 +++- .../llms/vertex_ai/batches/transformation.py | 9 +++- tests/unit/batches/test_batch_utils.py | 36 +++++++++++-- .../test_litellm_logging.py | 54 ++++++++++++++++++- .../custom_httpx/test_llm_http_handler.py | 4 +- tests/unit/test_cost_calculator.py | 24 +++++++++ 9 files changed, 173 insertions(+), 17 deletions(-) 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 35fc3510414..7ae03992007 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -712,7 +712,11 @@ class CheckBatchCost: from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials + from litellm.litellm_core_utils.litellm_logging import ( + deployment_pricing_model_info, + mask_api_base_credentials, + resolve_vertex_location_for_cost, + ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) @@ -849,6 +853,12 @@ class CheckBatchCost: custom_llm_provider=batch_file_provider, model_name=model_name, model_info=deployment_model_info, + vertex_location=resolve_vertex_location_for_cost( + custom_llm_provider=llm_provider, + litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True), + optional_params=None, + model=model_name, + ), ) error_file_failed_requests: Final = await count_error_file_failed_requests( response, diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 819a279a43c..cdc2625239f 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -73,6 +73,7 @@ async def calculate_batch_cost_and_usage( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, + vertex_location: str | None = None, ) -> BatchCostUsageResult: """ Calculate the cost and usage of a batch. @@ -82,16 +83,21 @@ async def calculate_batch_cost_and_usage( pricing. Threaded through to batch_cost_calculator so that deployment-specific pricing (e.g. input_cost_per_token_batches) is used instead of the global cost map. + vertex_location: The Vertex AI location the batch ran on; non-global + locations apply the model's regional-endpoint uplift multiplier. """ first_row: Final = file_content_dictionary[0] if file_content_dictionary else None if _uses_native_vertex_output(custom_llm_provider, model_name, first_row): - return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info) + return calculate_vertex_ai_batch_cost_and_usage( + file_content_dictionary, model_name, model_info=model_info, vertex_location=vertex_location + ) return _aggregate_batch_cost_usage_models( entries=file_content_dictionary, custom_llm_provider=custom_llm_provider, model_name=model_name, model_info=model_info, + vertex_location=vertex_location, ) @@ -101,6 +107,7 @@ async def _handle_completed_batch( model_name: str | None = None, litellm_params: dict | None = None, model_info: ModelInfo | None = None, + vertex_location: str | None = None, ) -> BatchCostUsageResult: """Fetch a completed batch's output file and aggregate its cost, usage, and models in a single pass over the JSONL lines, so the parsed file content is @@ -114,6 +121,8 @@ async def _handle_completed_batch( model_info: Optional deployment-level model info with custom pricing, threaded through so a deployment's configured rates win over the global cost map. + vertex_location: The Vertex AI location the batch ran on; non-global + locations apply the model's regional-endpoint uplift multiplier. """ # A completed batch whose request lines all failed has no output file - the # results are written to a separate error_file_id and output_file_id is None. @@ -141,7 +150,10 @@ async def _handle_completed_batch( output_file_result: Final = ( calculate_vertex_ai_batch_cost_and_usage( - _iter_batch_output_entries(file_content), model_name, model_info=model_info + _iter_batch_output_entries(file_content), + model_name, + model_info=model_info, + vertex_location=vertex_location, ) if _uses_native_vertex_output( custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None) @@ -151,6 +163,7 @@ async def _handle_completed_batch( custom_llm_provider=custom_llm_provider, model_name=model_name, model_info=model_info, + vertex_location=vertex_location, ) ) @@ -186,6 +199,7 @@ def _classify_output_line_stats( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, + vertex_location: str | None = None, ) -> Iterator[_BatchOutputLineStats | _LineOutcome]: """Classify every output line in a single pass, so counting failures never needs a second read of a potentially huge output file. A line the provider reported as @@ -196,7 +210,7 @@ def _classify_output_line_stats( if not _batch_response_was_successful(entry, custom_llm_provider): yield _LineOutcome.PROVIDER_FAILED continue - stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) + stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location) yield stats if stats is not None else _LineOutcome.UNCOSTABLE @@ -205,13 +219,14 @@ def _safe_output_line_stats( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, + vertex_location: str | None = None, ) -> _BatchOutputLineStats | None: """Return the stats for one provider-successful batch output line, or None when it cannot be costed, so a single bad line never aborts the whole batch's cost accounting.""" custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None try: - return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info) + return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location) except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch verbose_logger.warning( "batch output line could not be costed, so it is billed at $0 and the rest of the batch " @@ -227,6 +242,7 @@ def _compute_output_line_stats( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, + vertex_location: str | None = None, ) -> _BatchOutputLineStats: response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider) @@ -241,6 +257,7 @@ def _compute_output_line_stats( model_name=model_name, response_model=response_model, model_info=model_info, + vertex_location=vertex_location, ) return _BatchOutputLineStats( prompt_cost=line_prompt_cost, @@ -270,6 +287,7 @@ def _output_line_cost( model_name: str | None, response_model: str | None, model_info: ModelInfo | None, + vertex_location: str | None = None, ) -> tuple[float, float]: """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost @@ -290,6 +308,7 @@ def _output_line_cost( model=cost_model, custom_llm_provider=custom_llm_provider, model_info=model_info, + vertex_location=vertex_location, ) @@ -298,11 +317,14 @@ def _aggregate_batch_cost_usage_models( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, + vertex_location: str | None = None, ) -> BatchCostUsageResult: """Aggregate cost, usage, models, and pass/fail counts from batch output entries in a single pass, holding one small stats record per line instead of the parsed file.""" - all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info)) + all_results: Final = tuple( + _classify_output_line_stats(entries, custom_llm_provider, model_name, model_info, vertex_location) + ) line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats)) failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED) successful_requests: Final = len(all_results) - failed_requests @@ -349,6 +371,7 @@ def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: Iterable[dict], model_name: str | None = None, model_info: ModelInfo | None = None, + vertex_location: str | None = None, ) -> BatchCostUsageResult: """ Cost and usage of a native Vertex predictions.jsonl, one @@ -367,6 +390,7 @@ def calculate_vertex_ai_batch_cost_and_usage( model_info=model_info, calculate_usage=VertexGeminiConfig._calculate_usage, cost_calculator=batch_cost_calculator, + vertex_location=vertex_location, ) for row in vertex_ai_batch_responses ) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 845c62d91d6..80e96b8e888 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -36,6 +36,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( get_batch_cost_rates, get_billable_input_tokens, get_token_type_cost_breakdown, + get_vertex_regional_endpoint_uplift, parse_prompt_tokens_details, select_cost_metric_for_model, ) @@ -627,6 +628,7 @@ def cost_per_token( model=model, custom_llm_provider=custom_llm_provider, data_residency=data_residency, + vertex_location=vertex_location, ) elif call_type == "atranscription" or call_type == "transcription": if _transcription_usage_has_token_details(usage_block): @@ -2571,6 +2573,8 @@ def batch_cost_calculator( custom_llm_provider: str | None = None, model_info: ModelInfo | None = None, data_residency: str | None = None, + *, + vertex_location: str | None = None, ) -> tuple[float, float]: """ Calculate the cost of a batch job. @@ -2656,7 +2660,8 @@ def batch_cost_calculator( total_prompt_cost *= uplift total_completion_cost *= uplift - return total_prompt_cost, total_completion_cost + vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location) + return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift def _batch_prompt_cost( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index efa6fd5f1e8..eba580561d3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -485,7 +485,7 @@ def _published_pricing(deployment_model: str | None) -> ModelInfo | None: return None -def _resolve_vertex_location_for_cost( +def resolve_vertex_location_for_cost( custom_llm_provider: str | None, litellm_params: Mapping[str, object] | None, optional_params: Mapping[str, object] | None, @@ -1909,7 +1909,7 @@ class Logging(LiteLLMLoggingBaseClass): if hasattr(self, "litellm_params") and self.litellm_params else None ), - "vertex_location": _resolve_vertex_location_for_cost( + "vertex_location": resolve_vertex_location_for_cost( custom_llm_provider=self.model_call_details.get("custom_llm_provider", None), litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None), optional_params=self.optional_params, @@ -3267,6 +3267,12 @@ class Logging(LiteLLMLoggingBaseClass): model_name=self.get_deployment_model_for_cost(), litellm_params=self.litellm_params, model_info=self.get_router_deployment_model_info(), + vertex_location=resolve_vertex_location_for_cost( + custom_llm_provider=self.custom_llm_provider, + litellm_params=self.litellm_params if hasattr(self, "litellm_params") else None, + optional_params=None, + model=self.get_deployment_model_for_cost() or self.model, + ), ) result._hidden_params["response_cost"] = batch_result.cost diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 11c99d130be..be66b395b1c 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -54,6 +54,8 @@ class NativeVertexBatchCostCalculator(Protocol): model: str, custom_llm_provider: str | None = None, model_info: ModelInfo | None = None, + *, + vertex_location: str | None = None, ) -> tuple[float, float]: ... @@ -90,6 +92,7 @@ def native_vertex_batch_row_stats( model_info: ModelInfo | None, calculate_usage: Callable[[GenerateContentResponseBody], Usage], cost_calculator: NativeVertexBatchCostCalculator, + vertex_location: str | None = None, ) -> NativeVertexBatchRowStats | None: """ Usage and cost of one native Vertex predictions.jsonl row, a @@ -121,7 +124,11 @@ def native_vertex_batch_row_stats( ) try: prompt_cost, completion_cost = cost_calculator( - usage=usage, model=model, custom_llm_provider="vertex_ai", model_info=model_info + usage=usage, + model=model, + custom_llm_provider="vertex_ai", + model_info=model_info, + vertex_location=vertex_location, ) except Exception as e: # noqa: BLE001 # one unpriceable row must not abort the batch's cost accounting verbose_logger.warning( diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index dd95addac40..398b6fb1aed 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -640,7 +640,7 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model, model_info=None: bu.BatchCostUsageResult( + lambda content, model, model_info=None, vertex_location=None: bu.BatchCostUsageResult( cost=9.9, usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), models=["gemini-2.0-flash-001"], @@ -671,7 +671,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"), + lambda content, model, model_info=None, vertex_location=None: pytest.fail("raw vertex path should not run"), ) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") @@ -1112,6 +1112,36 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk assert result.failed_requests == 0 +@pytest.mark.asyncio +async def test_handle_completed_vertex_batch_regional_location_costs_uplift_over_global(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.files.main as files_main + + rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)] + + async def fake_afile_content(**kw): + return type("R", (), {"content": _vertex_jsonl(rows)})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + + global_result = await bu._handle_completed_batch( + _batch("gs://litellm-bucket/output/predictions.jsonl"), + custom_llm_provider="vertex_ai", + litellm_params={"vertex_project": "proj-1", "vertex_location": "global"}, + vertex_location="global", + ) + regional_result = await bu._handle_completed_batch( + _batch("gs://litellm-bucket/output/predictions.jsonl"), + custom_llm_provider="vertex_ai", + litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"}, + vertex_location="us-central1", + ) + + uplift = litellm.model_cost["vertex_ai/gemini-3.6-flash"]["regional_endpoint_uplift_multiplier"] + assert global_result.cost > 0 + assert regional_result.cost == pytest.approx(global_result.cost * uplift, rel=1e-9) + assert regional_result.usage == global_result.usage + + @pytest.mark.asyncio async def test_output_file_content_no_output_file_id_raises(): with pytest.raises(ValueError, match="Output file id is None"): @@ -1341,7 +1371,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) seen: dict = {} - def fake_vertex_calc(content, model, model_info=None): + def fake_vertex_calc(content, model, model_info=None, vertex_location=None): seen["content"] = content seen["model"] = model return bu.BatchCostUsageResult( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 10bc65201a1..c112a88177e 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1331,6 +1331,56 @@ class TestRetrieveBatchCostPassesModelIdentity: assert captured["model_info"] is not None assert captured["model_info"]["input_cost_per_token"] == 0.0 + @pytest.mark.asyncio + async def test_forwards_vertex_location_of_the_deployment(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils import litellm_logging as logging_module + + captured: dict[str, object] = {} + + from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.types.utils import LiteLLMBatch, Usage + + async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult: + captured.update(kwargs) + return BatchCostUsageResult( + cost=1.25, + usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), + models=["gemini-3.5-flash"], + successful_requests=1, + failed_requests=0, + ) + + monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch) + + obj = LitellmLogging( + model="vertex_ai/gemini-3.5-flash", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="batch-call-2", + function_id="f", + ) + obj.custom_llm_provider = "vertex_ai" + obj.litellm_params = {"vertex_project": "proj-1", "vertex_location": "us-central1"} + + batch = LiteLLMBatch( + id="batch_vertex", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id="file-out", + ) + + with contextlib.suppress(Exception): + await obj._async_success_handler_body(result=batch, start_time=None, end_time=None) + + assert captured, "handle_completed_batch was never called" + assert captured["vertex_location"] == "us-central1" + class TestRetrieveBatchPricesOnlyFinalBatches: """Regression (LIT-7048): retrieving a provider-id batch priced it on every poll. @@ -6796,10 +6846,10 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur def _resolve(custom_llm_provider, litellm_params, optional_params, model): from litellm.litellm_core_utils.litellm_logging import ( - _resolve_vertex_location_for_cost, + resolve_vertex_location_for_cost, ) - return _resolve_vertex_location_for_cost( + return resolve_vertex_location_for_cost( custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, optional_params=optional_params, diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 399e4dbf206..1dc5f717b04 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -2900,7 +2900,7 @@ async def test_async_anthropic_messages_handler_carries_deployment_vertex_locati from litellm.litellm_core_utils.litellm_logging import ( Logging, - _resolve_vertex_location_for_cost, + resolve_vertex_location_for_cost, ) monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5") @@ -2949,7 +2949,7 @@ async def test_async_anthropic_messages_handler_carries_deployment_vertex_locati global_deployment = await logging_obj_after_handler(GenericLiteLLMParams(vertex_location="global")) assert global_deployment.litellm_params["vertex_location"] == "global" assert ( - _resolve_vertex_location_for_cost( + resolve_vertex_location_for_cost( custom_llm_provider="vertex_ai", litellm_params=global_deployment.litellm_params, optional_params=global_deployment.optional_params, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 99dea6366f9..237740cf9e6 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -9,6 +9,7 @@ from pydantic import BaseModel import litellm from litellm.cost_calculator import ( + batch_cost_calculator, BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, @@ -3220,6 +3221,29 @@ def _batch_cache_usage() -> Usage: ) +def test_batch_cost_calculator_vertex_regional_costs_uplift_over_global(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Regression for https://github.com/BerriAI/litellm/issues/46880: a Vertex batch + on a regional location must cost the global batch times the model's + regional_endpoint_uplift_multiplier, the same invariant the online path keeps. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + usage = Usage(prompt_tokens=1_000_000, completion_tokens=1_000_000, total_tokens=2_000_000) + global_prompt, global_completion = batch_cost_calculator( + usage=usage, model="gemini-3.5-flash", custom_llm_provider="vertex_ai", vertex_location="global" + ) + regional_prompt, regional_completion = batch_cost_calculator( + usage=usage, model="gemini-3.5-flash", custom_llm_provider="vertex_ai", vertex_location="us-central1" + ) + + uplift = litellm.get_model_info("vertex_ai/gemini-3.5-flash")["regional_endpoint_uplift_multiplier"] + assert global_prompt + global_completion > 0 + assert regional_prompt == pytest.approx(global_prompt * uplift, rel=1e-9) + assert regional_completion == pytest.approx(global_completion * uplift, rel=1e-9) + + def test_batch_cost_calculator_prices_multimodal_tokens_at_modality_rates(): from litellm.cost_calculator import batch_cost_calculator