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 7449ec9882f..6205936e0ee 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t import asyncio import json +import re from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone @@ -39,6 +40,10 @@ if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch +def _error_names_resource(error: Exception, resource_id: str) -> bool: + return resource_id in re.findall(r"""[^\s"'`]+""", str(error)) + + CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( @@ -109,6 +114,10 @@ class CheckBatchCost: # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True self.batch_processed_support_confirmed: bool = False + self._poll_cursor: tuple[datetime, str] | None = None + self._poll_before: datetime | None = None + self._arrival_cursor: tuple[datetime, str] = (datetime.now(timezone.utc), "") + self._poll_arrivals: bool = False self._line_item_stores: Final[set[asyncio.Task[None]]] = set() self._line_item_store_turn: Final = asyncio.Lock() self._poll_idle: Final = asyncio.Event() @@ -161,29 +170,9 @@ class CheckBatchCost: 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 _managed_object_table(self.prisma_client).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) -> "Sequence[ManagedObjectRow]": """Query batch jobs without the batch_processed filter (for older schemas).""" - return await _managed_object_table(self.prisma_client).find_many( + return await self._find_jobs( where={ "file_purpose": "batch", "status": { @@ -197,10 +186,52 @@ class CheckBatchCost: ] }, }, - take=MAX_OBJECTS_PER_POLL_CYCLE, - order={"created_at": "asc"}, ) + async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]: + arrivals: Final = self._poll_arrivals + page: Final = await self._find_job_page(where, arrivals=arrivals) + self._poll_arrivals = not arrivals if page else arrivals + if page: + return page + return await self._find_job_page(where, arrivals=not arrivals) + + async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]: + cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor + if not arrivals and cursor is None: + self._poll_before = datetime.now(timezone.utc) + page: Final = await _managed_object_table(self.prisma_client).find_many( + where={ + **where, + "created_at": {"lte": datetime.now(timezone.utc) if arrivals else self._poll_before}, + **( + { + "OR": [ + {"created_at": {"gt": cursor[0]}}, + {"created_at": cursor[0], "id": {"gt": cursor[1]}}, + ] + } + if cursor is not None + else {} + ), + }, + take=MAX_OBJECTS_PER_POLL_CYCLE, + order=[{"created_at": "asc"}, {"id": "asc"}], + ) + if arrivals: + if page: + self._arrival_cursor = (page[-1].created_at, page[-1].id) + else: + self._poll_cursor = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None + return page + + def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]: + return { + "id": job.id, + "updated_at": job.updated_at, + **({"batch_processed": False} if self._has_batch_processed_column else {}), + } + async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None: """ Take a row that can never be costed out of the poll page. Leaving it selectable @@ -208,24 +239,19 @@ class CheckBatchCost: 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"} - ) + data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"} try: - await _managed_object_table(self.prisma_client).update( - where={"id": job.id}, + retired: Final = await _managed_object_table(self.prisma_client).update_many( + where=self._unprocessed_snapshot(job), data=data, ) except Exception as db_err: - verbose_proxy_logger.error( - f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}" - ) + verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}") + return + if retired == 0: return verbose_proxy_logger.warning( - f"CheckBatchCost: job {job.id} can never be costed ({reason}), " - "so it will no longer be polled" + f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled" ) async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool: @@ -312,13 +338,13 @@ class CheckBatchCost: from litellm.exceptions import NotFoundError - return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error) + return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id) 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.""" + 404 must not retire a potentially recoverable row.""" return self.llm_router.get_deployment(model_id=model_id) is not None @staticmethod @@ -333,7 +359,7 @@ class CheckBatchCost: if not output_file_id: return False - return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error) + return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id) async def _finalize_unbilled_terminal_job( self, job: "ManagedObjectRow", response: "LiteLLMBatch" @@ -355,18 +381,18 @@ class CheckBatchCost: db_batch_object=job, unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), ) - update_data: Final[dict] = { + update_data: Final[dict[str, object]] = { "status": response.status, "file_object": response.model_dump_json(), **({"batch_processed": True} if self._has_batch_processed_column else {}), } - await _managed_object_table(self.prisma_client).update( - where={"id": job.id}, + finalized: Final = await _managed_object_table(self.prisma_client).update_many( + where=self._unprocessed_snapshot(job), data=update_data, ) - verbose_proxy_logger.info( - f"CheckBatchCost: marked job {job.id} as {response.status} in DB" - ) + if finalized == 0: + return False + verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB") return True except Exception as db_err: verbose_proxy_logger.error( @@ -1047,6 +1073,7 @@ class CheckBatchCost: """ try: from litellm.integrations.prometheus import PrometheusLogger + prom_logger = PrometheusLogger.get_instance() except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}") @@ -1071,7 +1098,7 @@ class CheckBatchCost: # endpoint may transition a batch to "complete" before # CheckBatchCost runs. The batch_processed=False filter # already prevents reprocessing finished batches. - jobs = await _managed_object_table(self.prisma_client).find_many( + jobs = await self._find_jobs( where={ "file_purpose": "batch", "batch_processed": False, @@ -1084,8 +1111,6 @@ class CheckBatchCost: ] }, }, - take=MAX_OBJECTS_PER_POLL_CYCLE, - order={"created_at": "asc"}, ) self.batch_processed_support_confirmed = True except Exception as query_err: diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index 3c7735b53c0..a079fcf0b45 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -9,6 +9,7 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast from pydantic import TypeAdapter @@ -27,6 +28,12 @@ if TYPE_CHECKING: class ManagedObjectRow(Protocol): + @property + def created_at(self) -> datetime: ... + + @property + def updated_at(self) -> datetime: ... + @property def id(self) -> str: ... diff --git a/tests/integration/spend/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py index f68f10e3d8a..cdecabbbc9c 100644 --- a/tests/integration/spend/test_batch_poll_starvation.py +++ b/tests/integration/spend/test_batch_poll_starvation.py @@ -1,14 +1,19 @@ import json import os +from datetime import datetime, timezone +from functools import partial from hashlib import sha256 from typing import Final +from uuid import uuid4 +import httpx import pytest from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows -from integration._support.upstream import delete_scenario, register_scenario +from integration._support.database import read_rows, write_rows +from integration._support.prisma_clients import connected_proxy_prisma +from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE @@ -59,17 +64,17 @@ def _accepting_routes() -> dict[str, JsonResponse | TextResponse]: } -def _gone_at_provider_routes() -> RoutedResponse: +def _gone_at_provider_routes(*, retryable: bool = False) -> RoutedResponse: return RoutedResponse( content_type="application/x-routed", routes={ **_accepting_routes(), "GET /batches/batch-$REQUEST_ID": JsonResponse( content_type="application/json", - status=404, + status=503 if retryable else 404, body={ "error": { - "message": "No batch found with id 'batch-$REQUEST_ID'.", + "message": "Temporary failure" if retryable else "No batch found with id 'batch-$REQUEST_ID'.", "type": "invalid_request_error", "param": "id", "code": "batch_not_found", @@ -80,7 +85,7 @@ def _gone_at_provider_routes() -> RoutedResponse: ) -def _completed_routes() -> RoutedResponse: +def _completed_routes(*, ready: bool = True) -> RoutedResponse: output_line: Final = { "id": "batch_req_1", "custom_id": "r1", @@ -106,7 +111,8 @@ def _completed_routes() -> RoutedResponse: routes={ **_accepting_routes(), "GET /batches/batch-$REQUEST_ID": JsonResponse( - content_type="application/json", body=_batch("completed", "file-out-$REQUEST_ID") + content_type="application/json", + body=_batch("completed", "file-out-$REQUEST_ID") if ready else _batch("in_progress", None), ), "GET /files/file-out-$REQUEST_ID/content": TextResponse( content_type="application/jsonl", body=json.dumps(output_line, separators=(",", ":")) + "\n" @@ -115,7 +121,37 @@ def _completed_routes() -> RoutedResponse: ) -def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str: +def _batch_was_polled(upstream: ScenarioHandle) -> bool: + response: Final = httpx.get(f"{upstream.control_url}/__observations", timeout=15, trust_env=False) + response.raise_for_status() + observations: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]).validate_json(response.content) + return any( + item.get("method") == "GET" and string_value(item["path"]).startswith(f"/{upstream.scenario_id}/batches/") + for item in observations["requests"] + ) + + +def _missing_output_routes() -> RoutedResponse: + return RoutedResponse( + content_type="application/x-routed", + routes={ + **_completed_routes().routes, + "GET /files/file-out-$REQUEST_ID/content": JsonResponse( + content_type="application/json", + status=404, + body={ + "error": { + "message": "No file found with id 'file-out-$REQUEST_ID'.", + "type": "invalid_request_error", + "code": "file_not_found", + } + }, + ), + }, + ) + + +def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> tuple[str, ScenarioHandle]: scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}" handle: Final = register_scenario(scenario_id, routes) scenario.cleanups.callback(delete_scenario, handle) @@ -133,7 +169,7 @@ def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse }, ) scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) - return string_value(created["model_name"]) + return string_value(created["model_name"]), handle def _submitted_batch_id(gateway: Gateway, key: str, model_name: str) -> str: @@ -173,20 +209,43 @@ def _managed_rows(batch_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]: ) +@pytest.mark.parametrize("retryable", [False, True]) +@pytest.mark.parametrize("old_completed", [False, True]) @pytest.mark.timeout(180) @pytest.mark.covers("quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed") -def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling(gateway: Gateway) -> None: +def test_uncostable_batches_do_not_starve_billable_batches( + gateway: Gateway, retryable: bool, old_completed: bool +) -> None: with gateway.scenario() as scenario: - key: Final = scenario.key() + team_id: Final = scenario.team() + key: Final = scenario.key(team_id=team_id) gone_batch_ids: Final = tuple( _submitted_batch_id( - gateway, key, _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes()) + gateway, + key, + _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes(retryable=retryable))[0], ) for index in range(MAX_OBJECTS_PER_POLL_CYCLE) ) - costable_batch_id: Final = _submitted_batch_id( - gateway, key, _scripted_deployment(scenario, "costable", _completed_routes()) + missing_output_id: Final = _submitted_batch_id( + gateway, key, _scripted_deployment(scenario, "missing-output", _missing_output_routes())[0] ) + model, upstream = _scripted_deployment(scenario, "costable", _completed_routes(ready=False)) + costable_batch_id: Final = _submitted_batch_id(gateway, key, model) + for batch_id in (*gone_batch_ids, missing_output_id, costable_batch_id): + scenario.cleanups.callback( + write_rows, + 'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (batch_id,), + ) + if old_completed: + write_rows( + 'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s, ' + "created_at = NOW() - INTERVAL '30 days' WHERE unified_object_id = %s", + ("completed", costable_batch_id), + ) + assert eventually(lambda: _batch_was_polled(upstream), bool, seconds=120) + register_scenario(upstream.scenario_id, _completed_routes(), control_url=upstream.control_url) spend_rows: Final = eventually( lambda: read_rows( 'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" ' @@ -209,4 +268,78 @@ def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_pollin } ] assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}] - assert _managed_rows(gone_batch_ids) == [{"batch_processed": True}] * MAX_OBJECTS_PER_POLL_CYCLE + assert eventually( + lambda: _managed_rows((missing_output_id,)), + lambda rows: rows == [{"batch_processed": True}], + seconds=120, + ) == [{"batch_processed": True}] + expected_blockers: Final = [{"batch_processed": not retryable}] * MAX_OBJECTS_PER_POLL_CYCLE + blockers: Final = eventually( + lambda: _managed_rows(gone_batch_ids), lambda rows: rows == expected_blockers, seconds=120 + ) + assert blockers == expected_blockers + + expected_spend: Final = BATCH_COST_SHARE * ( + PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN + ) + for table, column, identity in ( + ("LiteLLM_VerificationToken", "token", sha256(key.encode()).hexdigest()), + ("LiteLLM_TeamTable", "team_id", team_id), + ): + balance: Final = eventually( + partial(read_rows, f'SELECT spend FROM "{table}" WHERE {column} = %s', (identity,)), + lambda rows: len(rows) == 1 and rows[0]["spend"] == pytest.approx(expected_spend), + seconds=70, + ) + assert balance == [{"spend": pytest.approx(expected_spend)}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("count", [2 * MAX_OBJECTS_PER_POLL_CYCLE, 2 * MAX_OBJECTS_PER_POLL_CYCLE + 1, 1000]) +async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) -> None: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + owner: Final = f"batch-scan-{uuid4().hex}" + + def insert_rows(group: str, size: int, created_at: datetime) -> tuple[str, ...]: + prefix: Final = f"{owner}-{group}-" + write_rows( + 'INSERT INTO "LiteLLM_ManagedObjectTable" ' + '(id, unified_object_id, model_object_id, file_object, file_purpose, created_by, created_at, updated_at) ' + "SELECT %s || lpad(n::text, 8, '0'), %s || n, %s || n, '{}'::jsonb, 'batch-scan-test', %s, " + "%s::timestamptz, NOW() FROM generate_series(0, %s::int - 1) n", + (prefix, prefix, prefix, owner, created_at.isoformat(), str(size)), + ) + return tuple(f"{prefix}{index:08}" for index in range(size)) + + try: + old_ids: Final = insert_rows("old", count, datetime(2000, 1, 1, tzinfo=timezone.utc)) + async with connected_proxy_prisma() as prisma: + logging: Final = ProxyLogging(UserApiKeyCache()) + router: Final = Router() + poller: Final = CheckBatchCost(logging, prisma, router) + where: Final = {"created_by": owner, "batch_processed": False} + first: Final = await poller._find_jobs(where) + assert tuple(row.id for row in first) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE] + + for start in range(MAX_OBJECTS_PER_POLL_CYCLE, count + 1, MAX_OBJECTS_PER_POLL_CYCLE): + new_ids: Final = insert_rows(str(start), MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc)) + arrivals: Final = await poller._find_jobs(where) + assert tuple(row.id for row in arrivals) == new_ids + retries: Final = await poller._find_jobs(where) + assert tuple(row.id for row in retries) == old_ids[start : start + MAX_OBJECTS_PER_POLL_CYCLE] + + revisited: Final = await poller._find_jobs(where) + assert tuple(row.id for row in revisited) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE] + latest_ids: Final = insert_rows("latest", MAX_OBJECTS_PER_POLL_CYCLE, datetime.now(timezone.utc)) + latest: Final = await poller._find_jobs(where) + assert tuple(row.id for row in latest) == latest_ids + + restarted: Final = CheckBatchCost(logging, prisma, router) + restart_page: Final = await restarted._find_jobs(where) + assert tuple(row.id for row in restart_page) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE] + finally: + write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)) diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index e3012dfcd13..241cf90eb5d 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -11,7 +11,8 @@ import json import logging from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from datetime import datetime +from copy import copy +from datetime import datetime, timezone from hashlib import sha256 from queue import SimpleQueue from types import SimpleNamespace @@ -217,7 +218,7 @@ class TestCheckBatchCost: find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE - assert find_call[1]["order"] == {"created_at": "asc"} + assert find_call[1]["order"] == [{"created_at": "asc"}, {"id": "asc"}] not_in = find_call[1]["where"]["status"]["not_in"] assert "stale_expired" in not_in # "complete"/"completed" are intentionally NOT excluded from the @@ -240,13 +241,13 @@ class TestCheckBatchCost: mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - side_effect=[Exception("column batch_processed does not exist"), []] + side_effect=[Exception("column batch_processed does not exist"), [], []] ) await check_batch_cost_instance.check_batch_cost() calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - assert len(calls) == 2 + assert len(calls) == 3 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where assert "stale_expired" in fallback_where["status"]["not_in"] @@ -267,8 +268,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - # Only one find_many call — the fallback directly, no primary query attempt - assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2 fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @@ -744,7 +744,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "Expected update() to be called exactly once for the completed job" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -877,7 +877,7 @@ class TestCheckBatchCost: mock_update_database.assert_awaited_once() assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -979,10 +979,10 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( f"Expected update() to be called exactly once for a {terminal_status} job" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True, ( "terminal-status update() must set batch_processed=True so polling stops" @@ -1085,9 +1085,13 @@ class TestCheckBatchCost: assert owner.user_id == "user-1" assert owner.team_id == "team-1" - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args - assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"} + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 + update_call = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args + assert update_call.kwargs["where"] == { + "id": "job-terminal-mint-1", + "batch_processed": False, + "updated_at": mock_job.updated_at, + } update_data = update_call.kwargs["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True @@ -1152,10 +1156,10 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "a completed batch with no output file must be marked processed exactly once" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == completed_status assert update_data["batch_processed"] is True, ( "completed-without-output update() must set batch_processed=True so polling stops" @@ -1371,7 +1375,7 @@ class TestCheckBatchCost: f"{terminal_status} batch with an output file must fetch results and be billed" ) mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == terminal_status, ( @@ -1544,10 +1548,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "a terminal batch with a 404ing output file must be retired, not retried forever" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1942,7 +1946,7 @@ class TestUnmanagedVertexRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 - assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2159,7 +2163,7 @@ class TestUnmanagedBedrockRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 - assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2740,9 +2744,10 @@ class TestPollPageStarvation: 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"} + prisma.db.litellm_managedobjecttable.update_many.assert_awaited() + call = prisma.db.litellm_managedobjecttable.update_many.call_args[1] + assert call["where"]["id"] == "job-no-model" + assert call["where"]["batch_processed"] is False assert call["data"] == {"batch_processed": True} @pytest.mark.asyncio @@ -2770,8 +2775,8 @@ class TestPollPageStarvation: 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} + prisma.db.litellm_managedobjecttable.update_many.assert_awaited() + assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2830,24 +2835,7 @@ class TestPollPageStarvation: 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} + assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_newer_batch_is_polled_once_dead_rows_are_retired(self): @@ -2885,7 +2873,11 @@ class TestPollPageStarvation: 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] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list + if "id" in call[1]["where"] + ] 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" @@ -2933,7 +2925,8 @@ class _FakeManagedObjectRow: self.team_id = None self.api_key = None self.request_tags = None - self.created_at = 1700000000 + self.created_at = datetime(9999, 1, 1, tzinfo=timezone.utc) + self.updated_at = datetime(2000, 1, 1, tzinfo=timezone.utc) self.file_object = json.dumps( { "id": "batch-456", @@ -2963,7 +2956,9 @@ class _FakeManagedObjectTable: def _matches(self, where: dict) -> bool: for key, value in where.items(): if key == "created_at": - return False + if "lt" in value and not self.row.created_at < value["lt"]: + return False + continue if key == "status": if self.row.status in value.get("not_in", []): return False @@ -2976,7 +2971,9 @@ class _FakeManagedObjectTable: async def _update_many(self, *, where: dict, data: dict) -> int: if not self._matches(where): return 0 - if "batch_processed" in where: + if "updated_at" in where: + self.journal.append("finalize") + elif "batch_processed" in where: self.journal.append("claim" if data.get("batch_processed") else "release") for key, value in data.items(): setattr(self.row, key, value) @@ -3224,6 +3221,89 @@ class TestMultiPodBatchCostClaim: ["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes ) + @pytest.mark.asyncio + @pytest.mark.parametrize("error_source", ("batch", "output")) + async def test_a_404_for_another_identifier_preserves_recovery(self, error_source: str) -> None: + import base64 + + import httpx + import litellm + import respx + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import SpecialEnums + + class SpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.costs: tuple[float, ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs["call_type"] == "aretrieve_batch": + cost: Final = kwargs["response_cost"] + assert isinstance(cost, float) + self.costs += (cost,) + + recorder: Final = SpendRecorder() + litellm.logging_callback_manager.add_litellm_async_success_callback(recorder) + litellm.register_model( + { + "batch-recovery-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": 0.01, + "output_cost_per_token_batches": 0.02, + } + } + ) + row: Final = _FakeManagedObjectRow() + row.unified_object_id = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode() + ) + .decode() + .rstrip("=") + ) + router: Final = self._router() + response: Final = router.aretrieve_batch.return_value + error: Final = litellm.NotFoundError( + message=f"No object found with id '{'batch-456' if error_source == 'batch' else _CLAIM_OUTPUT_FILE_ID}0'.", + model="gpt-4", + llm_provider="openai", + ) + router.aretrieve_batch.side_effect = (error, response) if error_source == "batch" else (response, response) + prisma: Final = self._prisma(row, []) + output: Final = json.dumps( + { + "response": { + "status_code": 200, + "body": { + "model": "batch-recovery-model", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + "error": None, + } + ).encode() + with respx.mock: + route: Final = respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content") + route.mock( + side_effect=error if error_source == "output" else None, + return_value=None if error_source == "output" else httpx.Response(200, content=output), + ) + await self._instance(prisma, router).check_batch_cost() + assert row.batch_processed is False + assert recorder.costs == () + route.mock(return_value=httpx.Response(200, content=output)) + await self._instance(prisma, router).check_batch_cost() + await self._instance(prisma, router).check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + assert recorder.costs == (pytest.approx(10 * 0.01 + 5 * 0.02),) + @staticmethod def _claim_calls(prisma) -> list: return [ @@ -3696,12 +3776,17 @@ class TestBatchLineItemsFromThePoller: assert logger.payloads == [] @pytest.mark.asyncio - async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch): + @pytest.mark.parametrize("database_error", [False, True]) + async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is( + self, monkeypatch: pytest.MonkeyPatch, database_error: bool + ) -> None: import litellm journal: list = [] row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1") - instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down")) + instance.prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + side_effect=RuntimeError("db down") if database_error else None, return_value=0 + ) logger = self._recording_logger(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) @@ -4111,6 +4196,127 @@ class TestCreatorTurnedOffMessageLogging: instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called() +@pytest.mark.asyncio +async def test_old_completed_batch_remains_available_for_costing() -> None: + row: Final = _FakeManagedObjectRow() + row.status = "completed" + row.created_at = datetime(2000, 1, 1, tzinfo=timezone.utc) + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + await poller._cleanup_stale_managed_objects() + + assert row.batch_processed is False + assert row.status == "completed" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("progress", ["claimed", "refreshed"]) +async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> None: + row: Final = _FakeManagedObjectRow() + selected: Final = copy(row) + if progress == "claimed": + row.batch_processed = True + row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc) + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + await poller._retire_job(selected, "provider batch disappeared") + + assert row.batch_processed is (progress == "claimed") + assert row.updated_at > selected.updated_at + + +@pytest.mark.asyncio +@pytest.mark.parametrize("progress", ["claimed", "refreshed"]) +async def test_stale_terminal_response_preserves_concurrent_progress(progress: str) -> None: + from litellm.types.utils import LiteLLMBatch + + row: Final = _FakeManagedObjectRow() + selected: Final = copy(row) + if progress == "claimed": + row.batch_processed = True + row.updated_at = datetime(2000, 1, 2, tzinfo=timezone.utc) + row.file_object = '{"revision":"fresh"}' + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + response: Final = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + status="failed", + ) + + assert await poller._finalize_unbilled_terminal_job(selected, response) is False + + assert row.file_object == '{"revision":"fresh"}' + assert row.status == "in_progress" + assert row.batch_processed is (progress == "claimed") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy", [False, True]) +async def test_permanently_lost_batch_retires_unchanged_row(legacy: bool) -> None: + row: Final = _FakeManagedObjectRow() + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + poller._has_batch_processed_column = not legacy + + await poller._retire_job(row, "provider batch disappeared") + + assert row.batch_processed is (not legacy) + assert row.status == ("stale_expired" if legacy else "in_progress") + + +@pytest.mark.asyncio +async def test_failed_retirement_keeps_batch_available() -> None: + row: Final = _FakeManagedObjectRow() + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + prisma.db.litellm_managedobjecttable.update_many = AsyncMock(side_effect=RuntimeError("database unavailable")) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + await poller._retire_job(row, "provider batch disappeared") + + assert row.batch_processed is False + assert row.status == "in_progress" + + +@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) +@pytest.mark.parametrize("other_id", ("{id}4", "prefix-{id}", "{id}/child", "{id}.backup", "{id}?generation=2")) +def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other_id: str) -> None: + from litellm.exceptions import NotFoundError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + error: Final = NotFoundError( + message=f"No object found with id '{other_id.format(id=resource_id)}'.", model="test", llm_provider="openai" + ) + assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is False + assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is False + + +@pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) +@pytest.mark.parametrize("message", ("No object found with id '{id}'.", 'File "{id}" does not exist', "{id} not found")) +def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None: + from litellm.exceptions import NotFoundError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + error: Final = NotFoundError(message=message.format(id=resource_id), model="test", llm_provider="openai") + assert CheckBatchCost._is_batch_gone_at_provider(error, resource_id) is True + assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True + + +@pytest.mark.parametrize("output_file_id", (None, "")) +def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None: + from litellm.exceptions import NotFoundError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + error: Final = NotFoundError(message="Output file not found", model="test", llm_provider="openai") + assert CheckBatchCost._is_output_file_gone_at_provider(error, output_file_id) is False + + @pytest.mark.asyncio @pytest.mark.parametrize( "kind", ["sync", "async", "custom", "overlapping", "general", "mixed", "sync_only", "executor_shutdown"]