From 1a5b941aeeb677eb77239f60a72ab2d258c8ed37 Mon Sep 17 00:00:00 2001 From: Charles Cheng Date: Wed, 12 Aug 2026 19:42:41 +0800 Subject: [PATCH] fix(check_batch_cost): retire permanently-unroutable and not-found batches _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. --- .../proxy/common_utils/check_batch_cost.py | 76 +++++++- .../proxy_unit_tests/test_check_batch_cost.py | 176 ++++++++++++++++++ 2 files changed, 248 insertions(+), 4 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 dc8f17fb665..ef3d0455069 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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}" diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6d7ada17ec5..57baf4ba8c9 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -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(