From ec52858865b8553fa2d7fad5cc2e701dc7ed8199 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Fri, 14 Aug 2026 00:20:10 -0400 Subject: [PATCH] fix(batches): only hand accounting to the poller once it can mark batches done The handoff asked whether the poller was running, when what matters is whether it will actually account for the batch. Those differ on a schema without the batch_processed column: the poller cannot filter on it, so it falls back to a query that excludes complete and completed rows, and it cannot set it either. A caller retrieving a provider-completed batch before the poller saw it therefore suppressed inline accounting, then marked the row complete, and the fallback query could never find it again. Nobody accounted for that batch, so its cost escaped the caller's budget entirely. The poller now publishes batch_processed_support_confirmed, set only once a filtered query has actually succeeded, and the handoff requires it. Defaulting to unconfirmed keeps accounting on the retrieve path in exactly the cases the poller would drop the batch, including the window before the poller's first cycle. All four combinations account exactly once: unconfirmed leaves the retrieve accounting and setting the marker, whether or not the column exists, and confirmed is only reachable when the column is present, where the poller accounts and sets it. A scheduler that hands back something other than a bound method leaves no poller to interrogate, which reads as unconfirmed rather than as working. --- .../proxy/common_utils/check_batch_cost.py | 2 + litellm/proxy/batches_endpoints/endpoints.py | 9 +- .../openai_files_endpoints/common_utils.py | 19 ++- .../proxy_unit_tests/test_check_batch_cost.py | 6 + .../test_files_common_utils.py | 130 +++++++++++++++++- 5 files changed, 152 insertions(+), 14 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 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()