fix(batches): only trust a 404 from the batch's own deployment

This commit is contained in:
mateo-berri 2026-08-12 21:19:37 -07:00
parent c11ebbed27
commit da84142288
2 changed files with 37 additions and 1 deletions

View file

@ -244,6 +244,13 @@ class CheckBatchCost:
return isinstance(error, (NotFoundError, openai.NotFoundError))
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
def _record_error(
prom_logger: Optional["PrometheusLogger"], error_type: str
@ -746,7 +753,7 @@ class CheckBatchCost:
)
if prom_logger:
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
if self._is_batch_gone_at_provider(e):
if self._is_batch_gone_at_provider(e) and self._batch_deployment_exists(model_id):
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
continue

View file

@ -1878,6 +1878,35 @@ class TestPollPageStarvation:
"batch_processed": True
}
@pytest.mark.asyncio
async def test_provider_404_with_deployment_gone_keeps_job(self):
"""With the batch's own deployment removed from the router, default fallbacks can
send the retrieve to a provider that never saw the batch. That 404 proves nothing,
so the row must stay unprocessed instead of losing its spend forever."""
import litellm
prisma = self._prisma(
[
self._job(
"job-misrouted",
self._encode("litellm_proxy;model_id:model-gone;llm_batch_id:batch_alive"),
)
]
)
llm_router = MagicMock()
llm_router.get_deployment = MagicMock(return_value=None)
llm_router.aretrieve_batch = AsyncMock(
side_effect=litellm.NotFoundError(
message="No batch found with id 'batch_alive'.",
model="model-gone",
llm_provider="openai",
)
)
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_transient_provider_error_keeps_job_for_retry(self):
"""A failure that may clear up (timeout, 5xx) must still leave the row unprocessed."""