scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

fix(check_batch_cost): retire permanently-unroutable and not-found batches

Browse source 
_resolve_job_routing() returned None for every routing failure and the
retrieval except-block continued on every provider error, so a row that
can never be costed (a unified id with no model_id, or a batch the
provider reports 404 for) stayed batch_processed=False forever. Since
the poll query orders oldest-first with a fixed page size, one such row
permanently occupies a slot and starves every batch behind it.

Only these two failures are provably permanent; every other routing/
retrieval failure is config-dependent (flag off, deployment missing,
transient error) and must keep retrying. Retire just the permanent ones
by marking batch_processed=True with a synthetic terminal status
(unroutable/not_found), excluded from both the primary and fallback
poll queries the same way the provider-native terminal statuses are.
This commit is contained in:
Charles Cheng 2026-08-12 19:42:41 +08:00
parent f64479e74d
commit 1a5b941aee
2 changed files with 248 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -11,6 +11,7 @@ from litellm.constants import (
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
)
from litellm.exceptions import NotFoundError as LiteLLMNotFoundError
if TYPE_CHECKING:
from litellm.integrations.prometheus import PrometheusLogger
@ -23,6 +24,20 @@ if TYPE_CHECKING:
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
# Statuses a row is written back with when a failure is provably permanent — no retry
# or config change can ever resolve it — so it must stop occupying the fixed-size,
# oldest-first poll window instead of being retried forever. Both are excluded from
# every find_many query the same way the provider-native terminal statuses are.
PERMANENTLY_UNPROCESSABLE_STATUSES = ("unroutable", "not_found")
class _UnroutableBatchError(Exception):
"""Raised by _resolve_job_routing for the one routing failure that is permanent:
a unified id that decodes successfully but embeds no model id. Every other
routing failure is config-dependent (enabling unmanaged tracking, restoring a
deployment) and must keep being retried, so it returns None instead of raising.
"""
class CheckBatchCost:
def __init__(
@ -136,7 +151,17 @@ class CheckBatchCost:
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
"status": {
"not_in": [
"completed",
"complete",
"failed",
"expired",
"cancelled",
"stale_expired",
*PERMANENTLY_UNPROCESSABLE_STATUSES,
]
},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
@ -160,6 +185,7 @@ class CheckBatchCost:
"complete",
"completed",
"stale_expired",
*PERMANENTLY_UNPROCESSABLE_STATUSES,
]
},
},
@ -174,6 +200,28 @@ class CheckBatchCost:
if prom_logger is not None:
prom_logger.record_check_batch_cost_error(error_type)
async def _retire_unprocessable_job(
self, job: "LiteLLM_ManagedObjectTable", status: str
) -> None:
"""
Permanently mark a row as processed without costing it, for a failure that is
provably permanent (see PERMANENTLY_UNPROCESSABLE_STATUSES). Leaving
batch_processed=False here would let the row re-occupy the fixed-size,
oldest-first poll window on every cycle and starve every batch behind it.
"""
update_data: Dict[str, Any] = {"status": status}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
try:
await self.prisma_client.db.litellm_managedobjecttable.update(
where={"id": job.id},
data=update_data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire job {job.id} (status={status}) in DB: {db_err}"
)
def _resolve_job_routing(
self,
job: "LiteLLM_ManagedObjectTable",
@ -187,7 +235,10 @@ class CheckBatchCost:
LiteLLM's own /v1/batches with a raw input_file_id) store the raw provider job id as
unified_object_id instead; when track_unmanaged_batch_cost is enabled the model is derived
from the provider-specific input_file_id layout (Vertex gs:// or Bedrock s3://) and mapped
to a matching deployment. Returns None (recording a metric) when the row can't be routed.
to a matching deployment. Returns None (recording a metric) when the row can't be routed
for a config-dependent reason that may resolve on a later cycle. Raises
_UnroutableBatchError for the one failure that is permanent: a unified id that
decodes but embeds no model id, which no retry or config change can fix.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
@ -204,7 +255,7 @@ class CheckBatchCost:
f"Skipping job {unified_object_id} because it is not a valid model id"
)
self._record_error(prom_logger, "invalid_model_id")
return None
raise _UnroutableBatchError(unified_object_id)
return model_id, get_batch_id_from_unified_batch_id(decoded)
if self._track_unmanaged_batch_cost:
@ -625,6 +676,7 @@ class CheckBatchCost:
"expired",
"cancelled",
"stale_expired",
*PERMANENTLY_UNPROCESSABLE_STATUSES,
]
},
},
@ -643,7 +695,15 @@ class CheckBatchCost:
else:
jobs = await self._fallback_find_jobs()
for job in jobs:
routing = self._resolve_job_routing(job, prom_logger)
try:
routing = self._resolve_job_routing(job, prom_logger)
except _UnroutableBatchError:
verbose_proxy_logger.info(
f"Retiring job {job.unified_object_id}: its unified id decodes but "
"embeds no model id, so it can never be routed"
)
await self._retire_unprocessable_job(job, status="unroutable")
continue
if routing is None:
continue
model_id, batch_id = routing
@ -661,6 +721,14 @@ class CheckBatchCost:
"batch_ignore_default_logging": True,
},
)
except LiteLLMNotFoundError as e:
verbose_proxy_logger.info(
f"Retiring job {job.unified_object_id}: provider reports batch "
f"{batch_id} not found — this is permanent, no retry will recover it: {e}"
)
self._record_error(prom_logger, "provider_batch_not_found")
await self._retire_unprocessable_job(job, status="not_found")
continue
except Exception as e:
verbose_proxy_logger.info(
f"Skipping job {job.unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}"

176
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -614,6 +614,182 @@ class TestCheckBatchCost:
mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0
), "a failed cost tracking attempt must not mark the job processed"
@pytest.mark.asyncio
async def test_decoded_id_missing_model_id_is_permanently_retired(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""A unified id that decodes successfully but embeds no model id can never be
routed, regardless of retries or config changes. Leaving batch_processed=False
would let it re-occupy the oldest-N poll window every cycle and starve every
batch behind it, so it must be retired (status=unroutable, batch_processed=True)
instead of just skipped.
"""
import base64
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_job = MagicMock()
mock_job.id = "job-unroutable-1"
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;llm_batch_id:batch-456"
).decode()
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
await check_batch_cost_instance.check_batch_cost()
mock_llm_router.aretrieve_batch.assert_not_called()
assert (
mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
), "an unroutable row must be retired, not left for the next cycle"
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[
1
]["data"]
assert update_data["status"] == "unroutable"
assert update_data["batch_processed"] is True
@pytest.mark.asyncio
async def test_invalid_unified_id_is_left_for_retry_not_retired(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""Unlike a decoded id with no model id, an id that isn't a recognized unified
id at all (e.g. unmanaged-batch tracking is off) is a config-dependent failure —
enabling the flag later would make it routable. It must keep being retried, not
be retired.
"""
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_job = MagicMock()
mock_job.id = "job-invalid-unified-1"
mock_job.unified_object_id = "some-raw-id"
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
with patch(_IS_B64, return_value=False):
await check_batch_cost_instance.check_batch_cost()
mock_llm_router.aretrieve_batch.assert_not_called()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0
@pytest.mark.asyncio
async def test_provider_not_found_error_retires_job(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""A provider 404 (the batch record is gone/expired on the provider's side) is
permanent — no retry will ever bring it back. It must be retired (status=
not_found, batch_processed=True) instead of occupying the poll window forever.
"""
import base64
from litellm.exceptions import NotFoundError
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_job = MagicMock()
mock_job.id = "job-404-1"
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_llm_router.aretrieve_batch = AsyncMock(
side_effect=NotFoundError(
message="batch_not_found", model="gpt-4", llm_provider="openai"
)
)
await check_batch_cost_instance.check_batch_cost()
assert (
mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
), "a provider 404 must retire the row, not leave it for the next cycle"
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[
1
]["data"]
assert update_data["status"] == "not_found"
assert update_data["batch_processed"] is True
@pytest.mark.asyncio
async def test_generic_provider_retrieval_error_leaves_job_unprocessed(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
"""A transient provider error (rate limit, timeout, 5xx) must NOT retire the
row — only a provably permanent failure (404) does. The row is left unprocessed
so the next poll retries it.
"""
import base64
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_job = MagicMock()
mock_job.id = "job-transient-1"
mock_job.unified_object_id = base64.urlsafe_b64encode(
b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456"
).decode()
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
mock_llm_router.aretrieve_batch = AsyncMock(
side_effect=Exception("temporary network blip")
)
await check_batch_cost_instance.check_batch_cost()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0
@pytest.mark.asyncio
async def test_primary_and_fallback_queries_exclude_permanently_retired_statuses(
self, check_batch_cost_instance, mock_prisma_client
):
"""Both the primary (batch_processed-aware) and fallback queries must exclude
the synthetic terminal statuses used to retire permanently-unroutable/not-found
rows, so a retired row on an older schema (no batch_processed column) isn't
re-selected forever.
"""
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[]
)
await check_batch_cost_instance.check_batch_cost()
primary_not_in = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[
1
]["where"]["status"]["not_in"]
assert "unroutable" in primary_not_in
assert "not_found" in primary_not_in
check_batch_cost_instance._has_batch_processed_column = False
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[]
)
await check_batch_cost_instance.check_batch_cost()
fallback_not_in = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[
1
]["where"]["status"]["not_in"]
assert "unroutable" in fallback_not_in
assert "not_found" in fallback_not_in
@pytest.mark.asyncio
@pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"])
async def test_terminal_status_marks_job_processed(