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 dc8f17fb665f..6fe37f0aacbc 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,7 +3,7 @@ """ from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -23,6 +23,15 @@ CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" +TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( + "completed", + "complete", + "failed", + "expired", + "cancelled", + "stale_expired", +) + class CheckBatchCost: def __init__( @@ -132,11 +141,11 @@ async def _cleanup_stale_managed_objects(self) -> None: in non-terminal states as 'stale_expired'. These will never complete and should not be polled. """ - cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) - result = await self.prisma_client.db.litellm_managedobjecttable.update_many( + cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) + result: Final = 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": list(TERMINAL_MANAGED_OBJECT_STATUSES)}, "created_at": {"lt": cutoff}, }, data={"status": "stale_expired"}, @@ -147,6 +156,26 @@ async def _cleanup_stale_managed_objects(self) -> None: f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" ) + if not self._has_batch_processed_column: + return + + # A row already in a terminal status is never rewritten by the sweep above, so + # without this it keeps a poll-page slot forever and starves newer batches. + retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many( + where={ + "file_purpose": "batch", + "batch_processed": False, + "status": {"in": ["complete", "completed"]}, + "created_at": {"lt": cutoff}, + }, + data={"batch_processed": True}, + ) + if retired > 0: + verbose_proxy_logger.warning( + f"CheckBatchCost: gave up on {retired} completed managed objects older than " + f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed" + ) + async def _fallback_find_jobs(self) -> list: """Query batch jobs without the batch_processed filter (for older schemas).""" return await self.prisma_client.db.litellm_managedobjecttable.find_many( @@ -167,6 +196,68 @@ async def _fallback_find_jobs(self) -> list: order={"created_at": "asc"}, ) + async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None: + """ + Take a row that can never be costed out of the poll page. Leaving it selectable + would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and + once enough such rows accumulate no newer batch is ever reached. Older schemas + without batch_processed can only be excluded through the status filter. + """ + data: Final = ( + {"batch_processed": True} + if self._has_batch_processed_column + else {"status": "stale_expired"} + ) + try: + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data=data, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}" + ) + return + verbose_proxy_logger.warning( + f"CheckBatchCost: job {job.id} can never be costed ({reason}), " + "so it will no longer be polled" + ) + + @staticmethod + def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool: + """A unified id that decodes but carries no model_id can never be routed.""" + from litellm.proxy.openai_files_endpoints.common_utils import ( + convert_b64_uid_to_unified_uid, + get_model_id_from_unified_batch_id, + ) + + decoded: Final = convert_b64_uid_to_unified_uid(job.unified_object_id) + return ( + decoded != job.unified_object_id + and get_model_id_from_unified_batch_id(decoded) is None + ) + + @staticmethod + def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool: + """ + A 404 naming the batch means the provider dropped its record of it, so no later + retrieve can ever succeed. A 404 about anything else, a renamed Azure deployment + or a fallback deployment that never saw this batch, is still fixable in config, so + it keeps retrying. + """ + import openai + + from litellm.exceptions import NotFoundError + + return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error) + + 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 @@ -645,6 +736,8 @@ async def check_batch_cost(self): for job in jobs: routing = self._resolve_job_routing(job, prom_logger) if routing is None: + if self._has_unified_id_without_model(job): + await self._retire_job(job, "unified object id has no model id") continue model_id, batch_id = routing @@ -667,6 +760,8 @@ async def check_batch_cost(self): ) if prom_logger: prom_logger.record_check_batch_cost_error("provider_retrieval_error") + if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id): + await self._retire_job(job, f"batch {batch_id} no longer exists at the provider") continue ## RETRIEVE THE BATCH JOB OUTPUT FILE diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6d7ada17ec52..fa274324fd6a 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1791,3 +1791,241 @@ async def test_named_key_still_owns_the_alias(self): metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") assert metadata["user_api_key_alias"] == "prod-key" + + +class TestPollPageStarvation: + """LIT-5462 regression: a row that can never be costed used to keep its slot in the + MAX_OBJECTS_PER_POLL_CYCLE page forever, so once enough of them accumulated no newer + batch was ever polled or costed.""" + + def _instance(self, prisma, llm_router): + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook.return_value = None + return CheckBatchCost( + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma, + llm_router=llm_router, + ) + + def _prisma(self, jobs): + prisma = MagicMock() + prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + prisma.db.litellm_managedobjecttable.update = AsyncMock() + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=jobs) + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + return prisma + + def _job(self, job_id, unified_object_id): + job = MagicMock() + job.id = job_id + job.unified_object_id = unified_object_id + job.created_by = "user-1" + return job + + @staticmethod + def _encode(unified_id: str) -> str: + import base64 + + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + @pytest.mark.asyncio + async def test_unified_id_without_model_id_is_retired(self): + """A unified id that decodes but carries no model_id is unroutable no matter what + the config says, so it must leave the poll page instead of being retried forever.""" + prisma = self._prisma( + [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) + llm_router = MagicMock() + llm_router.aretrieve_batch = AsyncMock() + + await self._instance(prisma, llm_router).check_batch_cost() + + llm_router.aretrieve_batch.assert_not_awaited() + prisma.db.litellm_managedobjecttable.update.assert_awaited_once() + call = prisma.db.litellm_managedobjecttable.update.call_args[1] + assert call["where"] == {"id": "job-no-model"} + assert call["data"] == {"batch_processed": True} + + @pytest.mark.asyncio + async def test_provider_404_retires_job(self): + """The provider dropping its record of the batch is permanent: no later retrieve + can succeed, so the row must stop occupying a slot.""" + import litellm + + prisma = self._prisma( + [ + self._job( + "job-gone", + self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_deadbeef"), + ) + ] + ) + llm_router = MagicMock() + llm_router.aretrieve_batch = AsyncMock( + side_effect=litellm.NotFoundError( + message="No batch found with id 'batch_deadbeef'.", + model="model-123", + llm_provider="openai", + ) + ) + + await self._instance(prisma, llm_router).check_batch_cost() + + prisma.db.litellm_managedobjecttable.update.assert_awaited_once() + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "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.""" + prisma = self._prisma( + [ + self._job( + "job-flaky", + self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_flaky"), + ) + ] + ) + llm_router = MagicMock() + llm_router.aretrieve_batch = AsyncMock(side_effect=Exception("connection reset")) + + await self._instance(prisma, llm_router).check_batch_cost() + + prisma.db.litellm_managedobjecttable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_retirement_falls_back_to_status_without_batch_processed_column(self): + """Older schemas have no batch_processed column, so the only way to stop selecting + the row is the status filter the poll query already applies.""" + prisma = self._prisma( + [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) + instance = self._instance(prisma, MagicMock()) + instance._has_batch_processed_column = False + + await instance.check_batch_cost() + + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "status": "stale_expired" + } + + @pytest.mark.asyncio + async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): + """A row already in a terminal status is never rewritten by the staleness sweep, so + it needs its own bound or it starves newer batches indefinitely.""" + prisma = self._prisma([]) + + await self._instance(prisma, MagicMock()).check_batch_cost() + + calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list + assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep" + where = calls[1][1]["where"] + assert where["file_purpose"] == "batch" + assert where["batch_processed"] is False + assert where["status"] == {"in": ["complete", "completed"]} + assert "created_at" in where + assert calls[1][1]["data"] == {"batch_processed": True} + + @pytest.mark.asyncio + async def test_newer_batch_is_polled_once_dead_rows_are_retired(self): + """The end state the customer cares about: dead rows retire on the cycle they are + first seen, and the healthy batch behind them keeps getting polled.""" + dead_rows = [ + self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model")), + self._job( + "job-gone", + self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_deadbeef"), + ), + ] + live_row = self._job( + "job-live", + self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_live"), + ) + prisma = self._prisma(dead_rows + [live_row]) + + import litellm + + in_progress = MagicMock() + in_progress.status = "in_progress" + + async def _retrieve(model, batch_id, litellm_metadata): + if batch_id == "batch_deadbeef": + raise litellm.NotFoundError( + message=f"No batch found with id '{batch_id}'.", + model=model, + llm_provider="openai", + ) + return in_progress + + llm_router = MagicMock() + llm_router.aretrieve_batch = AsyncMock(side_effect=_retrieve) + + await self._instance(prisma, llm_router).check_batch_cost() + + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update.call_args_list + ] + assert retired == ["job-no-model", "job-gone"] + assert ( + llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" + ), "the newer healthy batch must still be polled in the same cycle" + + @pytest.mark.asyncio + async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): + """A 404 about something other than the batch, e.g. a renamed Azure deployment, is + fixable in config, so the row must survive to be costed after the fix.""" + import litellm + + prisma = self._prisma( + [ + self._job( + "job-bad-deployment", + self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_real"), + ) + ] + ) + llm_router = MagicMock() + llm_router.aretrieve_batch = AsyncMock( + side_effect=litellm.NotFoundError( + message="Error code: 404 - DeploymentNotFound", + model="model-123", + llm_provider="azure", + ) + ) + + await self._instance(prisma, llm_router).check_batch_cost() + + prisma.db.litellm_managedobjecttable.update.assert_not_awaited()