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 6fe37f0aacb..990964dc81f 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -51,6 +51,7 @@ class CheckBatchCost: # Cached after the first poll cycle. Once we know the column is absent we skip # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True + self.batch_processed_support_confirmed: bool = False async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: """ @@ -722,6 +723,7 @@ class CheckBatchCost: take=MAX_OBJECTS_PER_POLL_CYCLE, order={"created_at": "asc"}, ) + self.batch_processed_support_confirmed = True except Exception as query_err: if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower(): raise diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index d5bc4ac0116..563c380d34b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -498,10 +498,11 @@ async def retrieve_batch( ) if unified_batch_id and batch_cost_poller_is_active(): - data["litellm_metadata"] = { - **(data.get("litellm_metadata") or {}), - "batch_ignore_default_logging": True, - } + litellm_metadata = data.get("litellm_metadata") + if not isinstance(litellm_metadata, dict): + litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend + data["litellm_metadata"] = litellm_metadata + litellm_metadata["batch_ignore_default_logging"] = True # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b8c250e718b..012c0e6ea5b 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1233,11 +1233,16 @@ async def get_batch_from_database( def batch_cost_poller_is_active() -> bool: """ - Whether the CheckBatchCost poller is running and will therefore account for a - managed batch's cost itself. + Whether the CheckBatchCost poller will account for a managed batch's cost itself. - False whenever the poller cannot be relied on: polling disabled by config, or the - job absent from the scheduler because the enterprise import failed. + False whenever the poller cannot be relied on: polling disabled by config, the job + absent from the scheduler because the enterprise import failed, or the poller not + yet having confirmed that the batch_processed column exists. That last condition + matters because the poller needs the column both to find outstanding batches and to + mark them accounted; without it the poller falls back to a query that excludes + terminal statuses, so a batch the retrieve path has already marked complete becomes + invisible to it. Defaulting to False until the poller confirms support keeps the + retrieve path accounting in exactly the cases the poller would drop the batch. """ from litellm.constants import PROXY_BATCH_POLLING_ENABLED @@ -1249,7 +1254,11 @@ def batch_cost_poller_is_active() -> bool: scheduler = getattr(proxy_server_module, "scheduler", None) if scheduler is None: return False - return scheduler.get_job("check_batch_cost_job") is not None + job = scheduler.get_job("check_batch_cost_job") + if job is None: + return False + poller = getattr(getattr(job, "func", None), "__self__", None) + return getattr(poller, "batch_processed_support_confirmed", False) is True except Exception: # noqa: BLE001 # scheduler backends raise varied types from get_job; an unreadable scheduler means the poller cannot be relied on return False diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index fa274324fd6..ce03dd33f85 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -143,6 +143,11 @@ class TestCheckBatchCost: assert "complete" not in not_in assert "completed" not in not_in assert find_call[1]["where"]["batch_processed"] is False + # A successful filtered query is the only proof the column exists. The retrieve + # path reads this to decide whether handing accounting to the poller is safe: + # without the column the poller's fallback query excludes complete/completed, so + # a batch already marked complete would never be accounted by anyone. + assert check_batch_cost_instance.batch_processed_support_confirmed is True @pytest.mark.asyncio async def test_fallback_query_used_when_batch_processed_missing( @@ -171,6 +176,7 @@ class TestCheckBatchCost: assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE # Column absence is now cached — next call should go straight to fallback assert check_batch_cost_instance._has_batch_processed_column is False + assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio async def test_column_absence_cached_across_cycles( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 4268d9c5e3e..20858a2a15f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -106,19 +106,44 @@ class _FakeScheduler: return self._job +class _FakePoller: + def __init__(self, confirmed): + self.batch_processed_support_confirmed = confirmed + + def check_batch_cost(self): + return None + + +def _job_for(poller): + if poller is None: + return None + job = MagicMock() + job.func = poller.check_batch_cost + return job + + @pytest.mark.parametrize( "polling_enabled, job, expected", [ - (True, object(), True), + (True, _job_for(_FakePoller(confirmed=True)), True), + (True, _job_for(_FakePoller(confirmed=False)), False), (True, None, False), - (False, object(), False), + (False, _job_for(_FakePoller(confirmed=True)), False), + ], + ids=[ + "poller-running-and-column-confirmed", + "poller-running-but-column-unconfirmed", + "job-absent-enterprise-import-failed", + "polling-disabled-by-config", ], - ids=["poller-running", "job-absent-enterprise-import-failed", "polling-disabled-by-config"], ) def test_batch_cost_poller_is_active(monkeypatch, polling_enabled, job, expected): """The predicate must only claim the poller when it can actually be relied on, so a - proxy with polling switched off or without the enterprise job keeps accounting for - batch cost on the retrieve path.""" + proxy with polling switched off, without the enterprise job, or whose poller has not + confirmed batch_processed support keeps accounting for batch cost on the retrieve + path. The unconfirmed case is the one that matters for legacy schemas: without the + column the poller falls back to a query excluding terminal statuses, so a batch the + retrieve path already marked complete would never be accounted by anyone.""" import litellm.constants import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -206,3 +231,98 @@ async def test_retrieving_a_completed_batch_still_marks_processed_without_a_cost assert data["batch_processed"] is True assert data["status"] == "complete" + + +def test_batch_cost_poller_is_active_is_false_when_the_job_has_no_bound_poller(monkeypatch): + """A scheduler that hands back a plain function rather than a bound method leaves no + poller to interrogate, so the predicate stays conservative instead of assuming the + column is supported.""" + import litellm.constants + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_cost_poller_is_active, + ) + + def unbound_check_batch_cost(): + return None + + job = MagicMock() + job.func = unbound_check_batch_cost + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", _FakeScheduler(job), raising=False) + + assert batch_cost_poller_is_active() is False + + +def test_batch_cost_poller_is_active_is_false_when_get_job_raises(monkeypatch): + """Scheduler backends raise varied types; an unreadable scheduler must not be read as + a working poller.""" + import litellm.constants + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_cost_poller_is_active, + ) + + class _ExplodingScheduler: + def get_job(self, job_id): + raise RuntimeError("scheduler not started") + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", _ExplodingScheduler(), raising=False) + + assert batch_cost_poller_is_active() is False + + + +@pytest.mark.asyncio +async def test_retrieving_a_batch_whose_status_is_unchanged_writes_nothing(monkeypatch): + """A caller polling an already-complete batch must not write at all, so repeated polls + cannot flip batch_processed or disturb whichever component owns accounting.""" + import litellm.proxy.openai_files_endpoints.common_utils as cu + + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: False) + monkeypatch.setattr(cu, "ensure_batch_response_managed_file_ids", AsyncMock()) + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + + db_batch_object = MagicMock() + db_batch_object.status = "completed" + + await cu.update_batch_in_database( + batch_id="unified-batch-id", + unified_batch_id="unified-batch-id", + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + db_batch_object=db_batch_object, + operation="retrieve", + ) + + update_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_batch_in_database_is_a_noop_for_unmanaged_batches(monkeypatch): + """Batches with no managed object row have neither the flag nor a poller queue entry, so + this path must leave them alone entirely.""" + import litellm.proxy.openai_files_endpoints.common_utils as cu + + prisma_client = MagicMock() + update_mock = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update = update_mock + + await cu.update_batch_in_database( + batch_id="batch-raw-xyz", + unified_batch_id=False, + response=_completed_batch(), + managed_files_obj=MagicMock(), + prisma_client=prisma_client, + verbose_proxy_logger=MagicMock(), + operation="retrieve", + ) + + update_mock.assert_not_awaited()