scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 199

fix(batches): apply the Vertex regional endpoint uplift to batch cost (#47466)

Browse source 
* 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 c6c303a618 (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>
This commit is contained in:
Kerry Lu 2026-10-19 18:51:51 -07:00 • committed by devin-ai-integration[bot]
parent 96684a2570
commit f585b558e9
9 changed files with 173 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -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,

34
litellm/batches/batch_utils.py
View file

@ -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
)

7
litellm/cost_calculator.py
View file

@ -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(

10
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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

9
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -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(

36
tests/unit/batches/test_batch_utils.py
View file

@ -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(

54
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -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,

4
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -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,

24
tests/unit/test_cost_calculator.py
View file

@ -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