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 1e58de0146c8..25b005973555 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -52,6 +52,33 @@ def __init__( # 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 + + @staticmethod + def _is_missing_batch_processed_column_error(err: Exception) -> bool: + message: Final = str(err).lower() + return "batch_processed" in message or "unknown column" in message or "does not exist" in message + + async def confirm_batch_processed_support(self) -> None: + """ + Probe the batch_processed column before the proxy serves traffic, so the retrieve + path never sees an unconfirmed poller on a schema that has the column and accounts + inline for a batch the first poll cycle then accounts again. + """ + try: + await self.prisma_client.db.litellm_managedobjecttable.find_first( + where={"file_purpose": "batch", "batch_processed": False} + ) + except Exception as probe_err: + if not self._is_missing_batch_processed_column_error(probe_err): + verbose_proxy_logger.debug( + f"CheckBatchCost: batch_processed probe failed, the poll cycle will confirm support: {probe_err}" + ) + return + self._has_batch_processed_column = False + verbose_proxy_logger.warning("CheckBatchCost: batch_processed column not found, querying without it") + return + self.batch_processed_support_confirmed = True async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: """ @@ -724,8 +751,9 @@ async def check_batch_cost(self): 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(): + if not self._is_missing_batch_processed_column_error(query_err): raise # Permanent schema gap — cache the result so future cycles skip straight to fallback self._has_batch_processed_column = False diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 37d267fcd6ea..f1b4c6b5b17f 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -54,6 +54,9 @@ normalize_mime_type_for_provider, resolve_managed_output_file_model_name, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( + request_tags_from_metadata, +) from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue] AllMessageValues, AsyncCursorPage, @@ -1146,6 +1149,7 @@ async def async_post_call_success_hook( ## Check if unified_file_id is in the response unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id + is_batch_create: Final = unified_file_id is not None model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) @@ -1216,6 +1220,7 @@ async def async_post_call_success_hook( model_mappings={model_id: provider_file_id}, user_api_key_dict=user_api_key_dict, ) + request_metadata: Final = data.get("litellm_metadata") await self.store_unified_object_id( unified_object_id=response.id, file_object=response, @@ -1223,6 +1228,8 @@ async def async_post_call_success_hook( model_object_id=original_response_id, file_purpose="batch", user_api_key_dict=user_api_key_dict, + request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}), + persist_attribution=is_batch_create, ) # Only record batch creation metric on actual create (not retrieve/cancel). diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 87f9927b1913..6952c0c6f899 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -25,6 +25,7 @@ _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, + batch_cost_poller_is_active, decode_model_from_file_id, encode_batch_response_ids, encode_file_id_with_model, @@ -497,6 +498,14 @@ async def retrieve_batch( "Batch %s is in non-terminal state %s, syncing with provider", batch_id, response.status ) + poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active() + if poller_owns_accounting: + 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 if model_from_id is not None: @@ -581,6 +590,7 @@ async def retrieve_batch( verbose_proxy_logger=verbose_proxy_logger, db_batch_object=db_batch_object, operation="retrieve", + poller_owns_accounting=poller_owns_accounting, ) ### CALL HOOKS ### - modify outgoing data diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f2e6fb633e14..b9af01e9aea2 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1256,6 +1256,38 @@ async def get_batch_from_database( return None, None +def batch_cost_poller_is_active() -> bool: + """ + 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, 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 + + if not PROXY_BATCH_POLLING_ENABLED: + return False + try: + import litellm.proxy.proxy_server as proxy_server_module + + scheduler = getattr(proxy_server_module, "scheduler", None) + if scheduler is None: + return False + 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 + + async def update_batch_in_database( batch_id: str, unified_batch_id: str | Literal[False], @@ -1266,6 +1298,7 @@ async def update_batch_in_database( db_batch_object=None, operation: str = "update", user_api_key_dict=None, + poller_owns_accounting: bool | None = None, ): """ Update batch status and object in ManagedObjectTable. @@ -1280,6 +1313,12 @@ async def update_batch_in_database( db_batch_object: Optional existing database object; fetched by unified_object_id when omitted operation: Description of operation ("update", "cancel", etc.) user_api_key_dict: Optional auth context for creating managed file IDs + poller_owns_accounting: Whether the caller already decided that the cost poller + owns this batch's accounting. Callers that suppress their own inline + accounting must pass the same decision they acted on, because re-deciding + here can observe a poller that became usable in between and leave the batch + unmarked after it was already accounted for, billing it twice. Left None by + callers that record no cost themselves. """ import litellm.utils @@ -1329,15 +1368,8 @@ async def update_batch_in_database( "updated_at": litellm.utils.get_utc_datetime(), } - # When a batch reaches completion, also mark batch_processed=True. - # The cost callback is enqueued asynchronously during the - # aretrieve_batch call that detected completion (via the @client - # decorator). It is not awaited, so there is a theoretical window - # where the callback hasn't executed yet. In practice the callback - # completes reliably. Setting the flag here unblocks file deletion - # which queries batch_processed=False. CheckBatchCost acts as a - # safety net for the rare case where the callback fails. - if db_status == "complete": + poller_owns: Final = batch_cost_poller_is_active() if poller_owns_accounting is None else poller_owns_accounting + if db_status == "complete" and not poller_owns: update_data["batch_processed"] = True try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 10ee2e10a39e..bda6fc254993 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8986,6 +8986,7 @@ async def _scheduled_ptu_rollup() -> None: llm_router=llm_router, track_unmanaged_batch_cost=general_settings.get("track_unmanaged_batch_cost", False), ) + await check_batch_cost_job.confirm_batch_processed_support() scheduler.add_job( check_batch_cost_job.check_batch_cost, "interval", diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6bb19a07d6c9..72f8b87dd16d 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -114,6 +114,45 @@ async def test_cleanup_scoped_to_batch_file_purpose( assert "stale_expired" in where["status"]["not_in"] assert "created_at" in where + @pytest.mark.asyncio + async def test_startup_probe_confirms_batch_processed_support( + self, check_batch_cost_instance, mock_prisma_client + ): + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + + await check_batch_cost_instance.confirm_batch_processed_support() + + probe_where = mock_prisma_client.db.litellm_managedobjecttable.find_first.call_args[1]["where"] + assert probe_where["batch_processed"] is False + assert check_batch_cost_instance.batch_processed_support_confirmed is True + assert check_batch_cost_instance._has_batch_processed_column is True + + @pytest.mark.asyncio + async def test_startup_probe_marks_column_absent( + self, check_batch_cost_instance, mock_prisma_client + ): + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( + side_effect=Exception("column batch_processed does not exist") + ) + + await check_batch_cost_instance.confirm_batch_processed_support() + + assert check_batch_cost_instance.batch_processed_support_confirmed is False + assert check_batch_cost_instance._has_batch_processed_column is False + + @pytest.mark.asyncio + async def test_startup_probe_transient_error_defers_to_poll_cycle( + self, check_batch_cost_instance, mock_prisma_client + ): + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( + side_effect=Exception("connection reset by peer") + ) + + await check_batch_cost_instance.confirm_batch_processed_support() + + assert check_batch_cost_instance.batch_processed_support_confirmed is False + assert check_batch_cost_instance._has_batch_processed_column is True + @pytest.mark.asyncio async def test_find_many_uses_pagination_and_excludes_stale( self, check_batch_cost_instance, mock_prisma_client @@ -143,6 +182,7 @@ async def test_find_many_uses_pagination_and_excludes_stale( assert "complete" not in not_in assert "completed" not in not_in assert find_call[1]["where"]["batch_processed"] is False + 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 +211,7 @@ async def test_fallback_query_used_when_batch_processed_missing( 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/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index d260f79a09ad..6cc31f991a3e 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -828,3 +828,49 @@ async def test_cost_job_and_retrieve_paths_mint_identical_unified_output_file_id model_id="model-deploy-xyz", model_name=cost_job_model_name, ) + + +@pytest.mark.asyncio +async def test_batch_create_hook_persists_creating_key_and_tags(): + """Regression: the /v1/batches create hook must persist the creating key and the + request's tags on the managed object row. CheckBatchCost, which owns the batch's + accounting once the retrieve path defers to it, bills whatever the row carries, and + without these columns the cost lands on the user alone and the key's spend and + budget never see it.""" + managed_files = _make_managed_files_instance() + creator = UserAPIKeyAuth(api_key="sk-the-creator", user_id="alice", parent_otel_span=None) + create_response = _make_batch_response(status="validating", output_file_id=None) + + await managed_files.async_post_call_success_hook( + data={"litellm_metadata": {"tags": ["env:prod", "team:ml"], "user_api_key": creator.api_key}}, + user_api_key_dict=creator, + response=create_response, + ) + + managed_files.store_unified_object_id.assert_awaited_once() + stored = managed_files.store_unified_object_id.await_args.kwargs + assert stored["persist_attribution"] is True + assert stored["request_tags"] == ("env:prod", "team:ml") + assert stored["user_api_key_dict"] is creator + + +@pytest.mark.asyncio +async def test_batch_retrieve_hook_does_not_claim_attribution(): + """A retrieve carries unified_batch_id but no unified_file_id, so it must not rewrite + the row's paying key to whoever happens to poll the batch.""" + managed_files = _make_managed_files_instance() + retrieve_response = _make_batch_response(status="in_progress", output_file_id=None) + retrieve_response._hidden_params = { + "unified_batch_id": "some-unified-batch-id", + "model_id": "model-deploy-xyz", + "model_name": "azure/gpt-4", + } + + await managed_files.async_post_call_success_hook( + data={"litellm_metadata": {"tags": ["poller:tag"]}}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None), + response=retrieve_response, + ) + + managed_files.store_unified_object_id.assert_awaited_once() + assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index aa5c63280b8b..824654dcf665 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2410,3 +2410,34 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc await call_cancel(cancel_harness, _unified_batch_id()) assert cancel_harness.router_acancel.call_count == 1 + + + + +@pytest.mark.asyncio +async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness): + with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router.aretrieve_batch.await_count == 1 + metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {} + assert metadata.get("batch_ignore_default_logging") is True + + +@pytest.mark.asyncio +async def test_retrieve__managed_batch_still_accounts_inline_without_a_poller(retrieve_harness): + with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=False)): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router.aretrieve_batch.await_count == 1 + metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {} + assert metadata.get("batch_ignore_default_logging") is None + + +@pytest.mark.asyncio +async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retrieve_harness): + with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): + await call_retrieve(retrieve_harness, "batch-raw-xyz") + + metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {} + assert metadata.get("batch_ignore_default_logging") is None 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 ad7f5e4725ae..6ffb7daaa2d5 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 @@ -97,6 +97,274 @@ def test_apply_unified_file_ids_swaps_all_three_ids(): ) +class _FakeScheduler: + def __init__(self, job): + self._job = job + + def get_job(self, job_id): + assert job_id == "check_batch_cost_job" + 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, _job_for(_FakePoller(confirmed=True)), True), + (True, _job_for(_FakePoller(confirmed=False)), False), + (True, None, 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", + ], +) +def test_batch_cost_poller_is_active(monkeypatch, polling_enabled, job, expected): + 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, + ) + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", polling_enabled, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", _FakeScheduler(job), raising=False) + + assert batch_cost_poller_is_active() is expected + + +def test_batch_cost_poller_is_active_is_false_when_no_scheduler_exists(monkeypatch): + 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, + ) + + monkeypatch.setattr(litellm.constants, "PROXY_BATCH_POLLING_ENABLED", True, raising=False) + monkeypatch.setattr(proxy_server_module, "scheduler", None, raising=False) + + assert batch_cost_poller_is_active() is False + + +def _completed_batch() -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-done", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id="file-out", + ) + + +async def _run_update(monkeypatch, poller_active: bool) -> dict: + import litellm.proxy.openai_files_endpoints.common_utils as cu + + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: poller_active) + 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 = "in_progress" + + 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", + ) + + assert update_mock.await_count == 1 + return update_mock.await_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_retrieving_a_completed_batch_leaves_batch_processed_to_the_cost_poller(monkeypatch): + data = await _run_update(monkeypatch, poller_active=True) + + assert "batch_processed" not in data + assert data["status"] == "complete" + + +@pytest.mark.asyncio +async def test_retrieving_a_completed_batch_still_marks_processed_without_a_cost_poller(monkeypatch): + data = await _run_update(monkeypatch, poller_active=False) + + 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): + 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): + 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): + 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): + 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() + + +@pytest.mark.asyncio +async def test_the_caller_s_accounting_decision_wins_over_a_later_poller_transition(monkeypatch): + import litellm.proxy.openai_files_endpoints.common_utils as cu + + monkeypatch.setattr(cu, "batch_cost_poller_is_active", lambda: True) + 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 = "in_progress" + + 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", + poller_owns_accounting=False, + ) + + data = update_mock.await_args.kwargs["data"] + assert data["batch_processed"] is True + assert data["status"] == "complete" + + +@pytest.mark.asyncio +async def test_a_caller_that_handed_off_accounting_still_leaves_the_marker_alone(monkeypatch): + 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 = "in_progress" + + 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", + poller_owns_accounting=True, + ) + + data = update_mock.await_args.kwargs["data"] + assert "batch_processed" not in data + assert data["status"] == "complete" + + # =========================================================================== # # add_internal_model_credentials - the snapshot that lets a completed # batch's output file be read, and therefore its cost be recorded diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 918d39646b0b..ba207242e295 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7070,6 +7070,44 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): assert ps.store_model_in_db is False +@pytest.mark.asyncio +async def test_batch_cost_poller_is_confirmed_before_serving(monkeypatch): + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.openai_files_endpoints.common_utils import batch_cost_poller_is_active + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.proxy_config", AsyncMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.PROXY_BATCH_POLLING_ENABLED", True), + patch("litellm.constants.PROXY_BATCH_POLLING_ENABLED", True), + patch("litellm.proxy.proxy_server.get_secret_bool", return_value=False), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + poller = proxy_server_module.scheduler.get_job("check_batch_cost_job").func.__self__ + assert poller.batch_processed_support_confirmed is True + assert batch_cost_poller_is_active() is True + probe_where = mock_prisma_client.db.litellm_managedobjecttable.find_first.call_args[1]["where"] + assert probe_where["batch_processed"] is False + + @pytest.mark.asyncio async def test_store_model_in_db_db_override_when_config_false(): """