ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(vertex-ai): fix zero cost/usage on completed Vertex AI batch jobs (#27912)

Browse Source 
* fix(vertex-ai): fix zero cost/usage on completed Vertex AI batch jobs

Vertex batch jobs recorded 0 spend and 0 tokens after PR #25627 added
automatic transformation of GCS predictions.jsonl to OpenAI format.

Two bugs fixed:

1. batch_utils.py: the Vertex-specific cost/usage reader
   (calculate_vertex_ai_batch_cost_and_usage) was always invoked and
   reads raw usageMetadata fields that no longer exist in the
   OpenAI-shaped output. Now the reader is only used when
   disable_vertex_batch_output_transformation=True; otherwise the
   generic path handles the already-transformed OpenAI-shaped content.

2. cost_calculator.py: batch_cost_calculator skipped the global
   litellm.get_model_info() lookup when a model_info dict was passed
   in, even when that dict had no pricing fields (e.g. deployment
   metadata with only id/db_model). It now falls back to the global
   pricing table when the provided model_info has no pricing data.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Update litellm/cost_calculator.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* fix(cost-calculator): use not-any guard for pricing fallback in batch_cost_calculator

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(cost-calculator): treat explicit zero batch pricing as set in model_info

The fallback to litellm.get_model_info() used truthy checks on pricing
fields, so 0.0 was treated as missing and replaced by global rates.
Use `is not None` like elsewhere in cost calculation. Add regression test.

Co-authored-by: Sameer Kankute <Sameerlite@users.noreply.github.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Sameer Kankute <Sameerlite@users.noreply.github.com>
This commit is contained in:
Sameer Kankute 2026-05-15 17:17:02 +05:30 committed by GitHub
parent cbdc70d544
commit c2efe9e422
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 223 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
litellm/batches/batch_utils.py
View File

@ -113,8 +113,11 @@ def _batch_cost_calculator(
"""
Calculate the cost of a batch based on the output file id
"""
# Handle Vertex AI with specialized method
if custom_llm_provider == "vertex_ai" and model_name:
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(
file_content_dictionary, model_name
)
@ -136,10 +139,13 @@ def calculate_vertex_ai_batch_cost_and_usage(
model_name: Optional[str] = None,
) -> Tuple[float, Usage]:
"""
Calculate both cost and usage from Vertex AI batch responses.
Calculate both cost and usage from raw Vertex AI batch responses.
Vertex AI batch output lines have format:
{"request": ..., "status": "", "response": {"candidates": [...], "usageMetadata": {...}}}
Used only when ``litellm.disable_vertex_batch_output_transformation = True``.
In that case the GCS predictions.jsonl is returned as-is, with each line in
the native Vertex format:
{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}}
usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
"""
@ -362,8 +368,11 @@ def _get_batch_job_total_usage_from_file_content(
"""
Get the tokens of a batch job from the file content
"""
# Handle Vertex AI with specialized method
if custom_llm_provider == "vertex_ai" and model_name:
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
_, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
file_content_dictionary, model_name
)

20
litellm/cost_calculator.py
View File

@ -2120,6 +2120,26 @@ def batch_cost_calculator(
)
except Exception:
model_info = None
elif not any(
model_info.get(k) is not None
for k in (
"input_cost_per_token_batches",
"input_cost_per_token",
"output_cost_per_token_batches",
"output_cost_per_token",
)
):
# model_info was provided (e.g. deployment metadata with only id/db_model)
# but carries no pricing fields. Fall back to the global pricing table so
# that standard model pricing is used instead of silently returning $0.
try:
global_info = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
if global_info:
model_info = global_info
except Exception:
pass
if not model_info:
return 0.0, 0.0

32
tests/batches_tests/test_batch_custom_pricing.py
View File

@ -8,6 +8,7 @@ are ignored by the batch cost pipeline because they are never threaded
through to `batch_cost_calculator`.
"""
import litellm
import pytest
from litellm.batches.batch_utils import (
@ -60,6 +61,37 @@ CUSTOM_MODEL_INFO = {
# --- tests ---
def test_batch_cost_calculator_explicit_zero_pricing_not_overridden_by_global(
monkeypatch,
):
"""
Explicit ``0`` / ``0.0`` pricing must count as present so we do not fall back
to the global pricing table (truthiness would treat zero as missing).
"""
usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
def fake_get_model_info(*args, **kwargs):
return {
"input_cost_per_token_batches": 1e-3,
"output_cost_per_token_batches": 2e-3,
}
monkeypatch.setattr(litellm, "get_model_info", fake_get_model_info)
prompt_cost, completion_cost = batch_cost_calculator(
usage=usage,
model="any-model",
custom_llm_provider="openai",
model_info={
"input_cost_per_token_batches": 0.0,
"output_cost_per_token_batches": 0.0,
},
)
assert prompt_cost == 0.0
assert completion_cost == 0.0
def test_batch_cost_calculator_uses_custom_model_info():
"""batch_cost_calculator should use model_info override when provided."""
usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)

155
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py
View File

@ -589,3 +589,158 @@ class TestVertexAIBatchCostCalculation:
assert usage.prompt_tokens == 0
assert usage.completion_tokens == 0
assert usage.total_tokens == 0
def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
"""
Regression test for the bug where Vertex batch cost/usage was always 0.
After PR #25627 (transform_file_content_response), the GCS predictions.jsonl
is rewritten into OpenAI batch shape before the cost-tracking path sees it.
With disable_vertex_batch_output_transformation=False (default), the content
is OpenAI-shaped, so _batch_cost_calculator must fall through to the generic
path rather than calling calculate_vertex_ai_batch_cost_and_usage (which only
reads raw usageMetadata fields).
"""
import litellm
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_get_batch_job_total_usage_from_file_content,
)
openai_shaped_responses = [
{
"id": "batch_req_abc123",
"custom_id": "request-1",
"response": {
"status_code": 200,
"request_id": "chatcmpl-xyz",
"body": {
"id": "chatcmpl-xyz",
"object": "chat.completion",
"model": "gemini-2.0-flash-001",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "Hello!"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
},
},
},
"error": None,
},
{
"id": "batch_req_def456",
"custom_id": "request-2",
"response": {
"status_code": 200,
"request_id": "chatcmpl-uvw",
"body": {
"id": "chatcmpl-uvw",
"object": "chat.completion",
"model": "gemini-2.0-flash-001",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "World!"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 8,
"completion_tokens": 3,
"total_tokens": 11,
},
},
},
"error": None,
},
]
original_flag = getattr(
litellm, "disable_vertex_batch_output_transformation", False
)
try:
litellm.disable_vertex_batch_output_transformation = False
cost = _batch_cost_calculator(
file_content_dictionary=openai_shaped_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
)
usage = _get_batch_job_total_usage_from_file_content(
file_content_dictionary=openai_shaped_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
)
finally:
litellm.disable_vertex_batch_output_transformation = original_flag
assert (
usage.prompt_tokens == 18
), f"expected 18 prompt tokens, got {usage.prompt_tokens}"
assert (
usage.completion_tokens == 8
), f"expected 8 completion tokens, got {usage.completion_tokens}"
assert (
usage.total_tokens == 26
), f"expected 26 total tokens, got {usage.total_tokens}"
assert (
cost > 0
), f"expected non-zero cost for completed Vertex batch, got {cost}"
def test_raw_vertex_output_still_works_when_transformation_disabled(self):
"""
When disable_vertex_batch_output_transformation=True the GCS file is returned
as raw Vertex predictions.jsonl; the specialized reader must be used.
"""
import litellm
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_get_batch_job_total_usage_from_file_content,
)
raw_vertex_responses = [
{
"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
"status": "",
"response": {
"candidates": [{"content": {"parts": [{"text": "Hello!"}]}}],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15,
},
},
"processed_time": "2026-01-01T00:00:00Z",
},
]
original_flag = getattr(
litellm, "disable_vertex_batch_output_transformation", False
)
try:
litellm.disable_vertex_batch_output_transformation = True
cost = _batch_cost_calculator(
file_content_dictionary=raw_vertex_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
)
usage = _get_batch_job_total_usage_from_file_content(
file_content_dictionary=raw_vertex_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
)
finally:
litellm.disable_vertex_batch_output_transformation = original_flag
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5
assert usage.total_tokens == 15
assert cost > 0, "raw Vertex shape should also produce non-zero cost"