diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 18849ef5b64..73ea8cf1991 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19622,7 +19622,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index ba5853b7c46..19e56e97338 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -7,8 +7,9 @@ import asyncio import os from collections.abc import Mapping +from datetime import datetime from types import MappingProxyType -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response from pydantic import TypeAdapter @@ -24,6 +25,7 @@ from litellm.proxy.batches_endpoints.litellm_executed_batches import ( LiteLLMExecutedBatchRunner, ManagedBatchStore, batch_error, + executed_batch_runner_lost, litellm_executed_provider_for, resolve_litellm_executed_provider, ) @@ -61,12 +63,15 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata from litellm.proxy.route_llm_request import raise_if_required_body_param_missing -from litellm.proxy.utils import ProxyLogging, handle_exception_on_proxy, is_known_model +from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository from litellm.router import Router from litellm.types.llms.openai import LiteLLMBatchCreateRequest from litellm.types.utils import LiteLLMBatch +if TYPE_CHECKING: + from prisma.models import LiteLLM_ManagedObjectTable + router: Final = APIRouter() _METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) @@ -79,7 +84,7 @@ def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None: def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyLogging) -> LiteLLMExecutedBatchRunner: - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import general_settings, prisma_client managed_files: Final = proxy_logging_obj.get_proxy_hook("managed_files") if prisma_client is None or not isinstance(managed_files, ManagedBatchStore): @@ -92,9 +97,36 @@ def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyL prisma_client=prisma_client, managed_files=managed_files, proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, ) +async def _batch_from_database( + batch_id: str, + unified_batch_id: str | Literal[False], + executed_batch: bool, + managed_files_obj: object, + prisma_client: PrismaClient | None, + llm_router: Router | None, + proxy_logging_obj: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, +) -> tuple["LiteLLM_ManagedObjectTable | None", LiteLLMBatch | None]: + row, batch = await get_batch_from_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + ) + updated_at: Final[object] = getattr(row, "updated_at", None) + if not executed_batch or batch is None or llm_router is None or not isinstance(updated_at, datetime): + return row, batch + if not executed_batch_runner_lost(batch.status, updated_at): + return row, batch + runner: Final = _litellm_executed_batch_runner(llm_router, proxy_logging_obj) + return row, await runner.fail_abandoned(batch, user_api_key_dict) + + async def _raise_when_input_file_must_be_managed(model: str, credentials: Mapping[str, object]) -> None: if await litellm_executed_provider_for(credentials) is None: return @@ -548,15 +580,18 @@ async def retrieve_batch( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") from litellm.proxy.proxy_server import prisma_client - db_batch_object, response = await get_batch_from_database( + executed_batch: Final = isinstance(unified_batch_id, str) and is_litellm_executed_batch(unified_batch_id) + db_batch_object, response = await _batch_from_database( batch_id=batch_id, unified_batch_id=unified_batch_id, + executed_batch=executed_batch, managed_files_obj=managed_files_obj, prisma_client=prisma_client, - verbose_proxy_logger=verbose_proxy_logger, + llm_router=llm_router, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, ) - executed_batch: Final = isinstance(unified_batch_id, str) and is_litellm_executed_batch(unified_batch_id) if executed_batch and response is None: raise batch_error(404, f"No batch found with id '{batch_id}'.") diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 642b2ea5f8c..7bd4c678183 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -3,6 +3,7 @@ import json import time from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timezone from itertools import pairwise from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable @@ -12,7 +13,7 @@ from openai.types.batch import Errors from openai.types.batch_error import BatchError from openai.types.batch_request_counts import BatchRequestCounts from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError -from typing_extensions import ReadOnly, TypedDict, assert_never +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -25,6 +26,7 @@ from litellm.llms.base_llm.files.storage_backend_factory import get_storage_back from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.models.managed_files import LiteLLM_ManagedFileTable from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.openai_files_endpoints.common_utils import ( LITELLM_EXECUTED_BATCH_ID_PREFIX, @@ -44,12 +46,26 @@ if TYPE_CHECKING: BatchEndpoint: TypeAlias = Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"] BatchStatus: TypeAlias = Literal["in_progress", "finalizing", "completed", "failed", "cancelling", "cancelled"] - TERMINAL_BATCH_STATUSES: Final[frozenset[str]] = frozenset({"completed", "failed", "cancelled", "expired"}) +_STOP_STATUSES: Final[frozenset[str]] = TERMINAL_BATCH_STATUSES | frozenset({"cancelling"}) _BATCH_ENDPOINT_ADAPTER: Final[TypeAdapter[BatchEndpoint]] = TypeAdapter(BatchEndpoint) _CANCEL_POLL_SECONDS: Final = 1.0 +_HEARTBEAT_SECONDS: Final = 30.0 +_STALE_AFTER_SECONDS: Final = 180.0 _FILES_API_PROBE_TIMEOUT_SECONDS: Final = 5.0 _COMPLETION_WINDOW_SECONDS: Final = 24 * 60 * 60 +_RUNNER_LOST_MESSAGE: Final = "the proxy replica running this batch stopped before it finished; resubmit the batch" +_ROUTER_METHODS: Final[Mapping[BatchEndpoint, str]] = MappingProxyType( + { + "/v1/chat/completions": "acompletion", + "/v1/completions": "atext_completion", + "/v1/embeddings": "aembedding", + "/v1/responses": "aresponses", + } +) +_CANCELLING_TRANSITIONS: Final[Mapping[BatchStatus, BatchStatus]] = MappingProxyType( + {"completed": "cancelled", "in_progress": "cancelling", "finalizing": "cancelling"} +) LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE: Final = ( "upload it through POST /v1/files with purpose=batch and either the x-litellm-model header or the " "target_model_names form field naming the model, so LiteLLM keeps the file and runs the batch itself" @@ -165,20 +181,6 @@ class _RouterCall(Protocol): def __call__(self, **params: object) -> Awaitable[object]: ... # kwargs-ok: the request body is passed as keywords -def _router_method_name(endpoint: BatchEndpoint) -> str: - match endpoint: - case "/v1/chat/completions": - return "acompletion" - case "/v1/completions": - return "atext_completion" - case "/v1/embeddings": - return "aembedding" - case "/v1/responses": - return "aresponses" - case _: - assert_never(endpoint) - - def litellm_executed_provider_of(credentials: Mapping[str, object]) -> str | None: explicit_provider: Final = credentials.get("custom_llm_provider") provider: Final = ( @@ -197,12 +199,20 @@ class FilesApiProbe(Protocol): async def __call__(self, api_base: str, api_key: str | None) -> bool: ... +class BodyRejection(Protocol): + def __call__(self, body: Mapping[str, object], /) -> str | None: ... + + async def upstream_lacks_files_api(api_base: str, api_key: str | None, http_client: _HttpGetter | None = None) -> bool: client: Final = http_client or get_async_httpx_client(llm_provider=LlmProviders.HOSTED_VLLM) try: response: Final = await client.get( f"{api_base.rstrip('/')}/files", - headers={"Authorization": f"Bearer {api_key}"} if api_key else None, + headers=( + {"Authorization": f"Bearer {api_key}"} # mutable-ok: AsyncHTTPHandler.get wants a plain dict + if api_key + else None + ), timeout=_FILES_API_PROBE_TIMEOUT_SECONDS, ) except httpx.HTTPError: @@ -266,7 +276,13 @@ def _validation_reason(error: ValidationError) -> str: ) -def _parse_line(line_number: int, raw: bytes, endpoint: BatchEndpoint) -> BatchInputLine | InvalidBatchInput: +def _accept_every_body(_body: Mapping[str, object]) -> str | None: + return None + + +def _parse_line( + line_number: int, raw: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection +) -> BatchInputLine | InvalidBatchInput: try: line: Final = BatchInputLine.model_validate_json(raw) except ValidationError as e: @@ -275,14 +291,19 @@ def _parse_line(line_number: int, raw: bytes, endpoint: BatchEndpoint) -> BatchI return InvalidBatchInput(line_number, f"url {line.url!r} does not match the batch endpoint {endpoint!r}") if line.body.get("stream"): return InvalidBatchInput(line_number, "streaming requests are not supported in a batch") + rejection: Final = reject_body(line.body) + if rejection is not None: + return InvalidBatchInput(line_number, rejection) return line -def parse_batch_input(content: bytes, endpoint: BatchEndpoint) -> tuple[BatchInputLine, ...] | InvalidBatchInput: +def parse_batch_input( + content: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection = _accept_every_body +) -> tuple[BatchInputLine, ...] | InvalidBatchInput: raw_lines: Final = tuple((number, raw) for number, raw in enumerate(content.splitlines(), start=1) if raw.strip()) if not raw_lines: return InvalidBatchInput(None, "the input file has no requests") - parsed: Final = tuple(_parse_line(number, raw, endpoint) for number, raw in raw_lines) + parsed: Final = tuple(_parse_line(number, raw, endpoint, reject_body) for number, raw in raw_lines) first_invalid: Final = next((item for item in parsed if isinstance(item, InvalidBatchInput)), None) if first_invalid is not None: return first_invalid @@ -345,37 +366,35 @@ def _dump(response: object) -> Mapping[str, object]: def _resolve_transition(current_status: str, requested: BatchStatus) -> BatchStatus: if current_status != "cancelling": return requested - match requested: - case "completed": - return "cancelled" - case "in_progress" | "finalizing": - return "cancelling" - case "failed" | "cancelling" | "cancelled": - return requested - case _: - assert_never(requested) + return _CANCELLING_TRANSITIONS.get(requested, requested) + + +def executed_batch_runner_lost(status: str, updated_at: datetime) -> bool: + if status in TERMINAL_BATCH_STATUSES: + return False + return (datetime.now(timezone.utc) - updated_at).total_seconds() > _STALE_AFTER_SECONDS def _llm_batch_id_of(unified_batch_id: str) -> str: return get_batch_id_from_unified_batch_id(convert_b64_uid_to_unified_uid(unified_batch_id)) -class _CancelWatch: +class _StopWatch: def __init__(self, load_status: Callable[[], Awaitable[str | None]], interval_seconds: float) -> None: self._load_status = load_status self._interval_seconds = interval_seconds self._checked_at = float("-inf") - self._cancelling = False + self._stopped = False - async def cancelling(self) -> bool: - if self._cancelling: + async def stopped(self) -> bool: + if self._stopped: return True now: Final = time.monotonic() if now - self._checked_at < self._interval_seconds: return False self._checked_at = now - self._cancelling = await self._load_status() == "cancelling" - return self._cancelling + self._stopped = await self._load_status() in _STOP_STATUSES + return self._stopped class LiteLLMExecutedBatchRunner: @@ -385,7 +404,9 @@ class LiteLLMExecutedBatchRunner: prisma_client: PrismaClient, managed_files: ManagedBatchStore, proxy_logging_obj: ProxyLogging, + general_settings: Mapping[str, object], concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY, + heartbeat_seconds: float = _HEARTBEAT_SECONDS, storage_backend_factory: _StorageBackendFactory = get_storage_backend, upload_result_file: _ResultFileUploader = StorageBackendFileService.upload_file_to_storage_backend, ) -> None: @@ -393,7 +414,9 @@ class LiteLLMExecutedBatchRunner: self.prisma_client = prisma_client self.managed_files = managed_files self.proxy_logging_obj = proxy_logging_obj + self.general_settings = general_settings self.concurrency = concurrency + self.heartbeat_seconds = heartbeat_seconds self.storage_backend_factory = storage_backend_factory self.upload_result_file = upload_result_file @@ -408,7 +431,7 @@ class LiteLLMExecutedBatchRunner: ) -> LiteLLMBatch: endpoint: Final = _validate_endpoint(create_request.get("endpoint")) content: Final = await self._download_input(unified_input_file_id, user_api_key_dict) - parsed: Final = parse_batch_input(content, endpoint) + parsed: Final = parse_batch_input(content, endpoint, self._body_rejection(model)) if isinstance(parsed, InvalidBatchInput): raise batch_error(400, f"Invalid batch input file: {parsed.describe()}") llm_batch_id: Final = f"{LITELLM_EXECUTED_BATCH_ID_PREFIX}{uuid_module.uuid4().hex}" @@ -468,6 +491,30 @@ class LiteLLMExecutedBatchRunner: await self._store(cancelling, user_api_key_dict) return cancelling + async def fail_abandoned(self, batch: LiteLLMBatch, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch: + error: Final = BatchError(message=_RUNNER_LOST_MESSAGE, code="runner_lost") + errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list + failed: Final = batch.model_copy( + update=MappingProxyType({"status": "failed", "failed_at": int(time.time()), "errors": errors}) + ) + await self._store(failed, user_api_key_dict) + return failed + + def _body_rejection(self, model: str) -> BodyRejection: + def reject(body: Mapping[str, object]) -> str | None: + try: + is_request_body_safe( + request_body=dict(body), # mutable-ok: is_request_body_safe takes a dict + general_settings=dict(self.general_settings), # mutable-ok: is_request_body_safe takes a dict + llm_router=self.llm_router, + model=model, + ) + except ValueError as e: + return str(e) + return None + + return reject + async def _download_input(self, unified_input_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bytes: stored: Final = await self.managed_files.get_unified_file_id( unified_input_file_id, litellm_parent_otel_span=user_api_key_dict.parent_otel_span @@ -485,6 +532,7 @@ class LiteLLMExecutedBatchRunner: raise batch_error(400, str(e)) async def _run(self, run: _BatchRun) -> None: + heartbeat: Final = asyncio.create_task(self._heartbeat(run)) try: await self._execute(run) except Exception as e: # noqa: BLE001 # whatever fails, the batch must end up marked failed @@ -497,14 +545,31 @@ class LiteLLMExecutedBatchRunner: verbose_proxy_logger.exception( "LiteLLM-executed batch %s could not be marked failed: %s", run.unified_batch_id, advance_error ) + finally: + heartbeat.cancel() + + async def _heartbeat(self, run: _BatchRun) -> None: + while True: + await asyncio.sleep(self.heartbeat_seconds) + try: + await self._touch(run) + except Exception as e: # noqa: BLE001 # a missed beat is logged and the next one retries + verbose_proxy_logger.warning("LiteLLM-executed batch %s heartbeat failed: %s", run.unified_batch_id, e) + + async def _touch(self, run: _BatchRun) -> None: + await ManagedObjectRepository(self.prisma_client).table.update_many( + where={"unified_object_id": run.unified_batch_id}, # mutable-ok: Prisma filter + data={"updated_by": run.user_api_key_dict.user_id}, # mutable-ok: Prisma payload + ) async def _execute(self, run: _BatchRun) -> None: await self._advance(run, "in_progress") - watch: Final = _CancelWatch(lambda: self._load_status(run.unified_batch_id), _CANCEL_POLL_SECONDS) + watch: Final = _StopWatch(lambda: self._load_status(run.unified_batch_id), _CANCEL_POLL_SECONDS) semaphore: Final = asyncio.Semaphore(self.concurrency) results: Final = await asyncio.gather(*(self._run_row(run, line, watch, semaphore) for line in run.lines)) outcomes: Final = tuple(outcome for outcome in results if outcome is not None) - await self._advance(run, "finalizing") + if await self._advance(run, "finalizing") is None: + return succeeded: Final = tuple(outcome for outcome in outcomes if outcome.succeeded) failed: Final = tuple(outcome for outcome in outcomes if not outcome.succeeded) output_file_id: Final = await self._upload_results(run, "output", succeeded) @@ -519,10 +584,10 @@ class LiteLLMExecutedBatchRunner: ) async def _run_row( - self, run: _BatchRun, line: BatchInputLine, watch: _CancelWatch, semaphore: asyncio.Semaphore + self, run: _BatchRun, line: BatchInputLine, watch: _StopWatch, semaphore: asyncio.Semaphore ) -> RowOutcome | None: async with semaphore: - if await watch.cancelling(): + if await watch.stopped(): return None try: body: Final = await self._dispatch(run, line) @@ -537,7 +602,7 @@ class LiteLLMExecutedBatchRunner: return _dump(await self._router_call(run.endpoint)(**params)) def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall: - method: Final[object] = getattr(self.llm_router, _router_method_name(endpoint), None) + method: Final[object] = getattr(self.llm_router, _ROUTER_METHODS[endpoint], None) if not isinstance(method, _RouterCall): raise TypeError(f"the router has no callable for {endpoint}") return method @@ -574,15 +639,20 @@ class LiteLLMExecutedBatchRunner: ) return file_object.id - async def _advance(self, run: _BatchRun, requested: BatchStatus, fields: Mapping[str, object] = _NO_FIELDS) -> None: + async def _advance( + self, run: _BatchRun, requested: BatchStatus, fields: Mapping[str, object] = _NO_FIELDS + ) -> BatchStatus | None: current: Final = await self._load_batch(run.unified_batch_id) if current is None: raise RuntimeError(f"Batch {run.unified_batch_id} is no longer stored") + if current.status in TERMINAL_BATCH_STATUSES: + return None status: Final = _resolve_transition(current.status, requested) updated: Final = current.model_copy( update=MappingProxyType({**fields, "status": status, f"{status}_at": int(time.time())}) ) await self._store(updated, run.user_api_key_dict) + return status async def _store(self, batch: LiteLLMBatch, user_api_key_dict: UserAPIKeyAuth) -> None: await self.managed_files.store_unified_object_id( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b8dcb89baef..7356d197be9 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -118,31 +118,29 @@ async def _litellm_executed_batch_input_model( executed: Final = tuple( candidate for candidate, provider in zip(candidates, providers, strict=True) if provider is not None ) - match executed: - case (): - return None - case _ if purpose != "batch": - raise ProxyException( - message=( - f"The server behind {', '.join(executed)} has no Files API, so LiteLLM keeps only batch input " - f"files for it and runs the batch itself: upload with purpose=batch; got purpose={purpose}" - ), - type="invalid_request_error", - param="purpose", - code=400, - ) - case (only,) if len(candidates) == 1: - return only - case _: - raise ProxyException( - message=( - f"LiteLLM runs batches for {', '.join(executed)} itself and keeps their input files, so a batch " - f"input file can target only that one model; got target_model_names={', '.join(candidates)}" - ), - type="invalid_request_error", - param="target_model_names", - code=400, - ) + if not executed: + return None + if purpose != "batch": + raise ProxyException( + message=( + f"The server behind {', '.join(executed)} has no Files API, so LiteLLM keeps only batch input " + f"files for it and runs the batch itself: upload with purpose=batch; got purpose={purpose}" + ), + type="invalid_request_error", + param="purpose", + code=400, + ) + if len(candidates) == 1: + return executed[0] + raise ProxyException( + message=( + f"LiteLLM runs batches for {', '.join(executed)} itself and keeps their input files, so a batch " + f"input file can target only that one model; got target_model_names={', '.join(candidates)}" + ), + type="invalid_request_error", + param="target_model_names", + code=400, + ) _MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None) diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index b4a36336c22..66dbcd87c0b 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -12,6 +12,7 @@ from typing import Any, Final, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid as uuid_module +from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -105,8 +106,9 @@ class StorageBackendFileService: storage_url=storage_url, ) - # Store in managed files if target_model_names provided - if target_model_names: + if not target_model_names: + return file_object + try: await StorageBackendFileService._store_in_managed_files( file_object=file_object, file_data=file_data, @@ -116,9 +118,25 @@ class StorageBackendFileService: proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, ) - + except Exception: + await StorageBackendFileService._discard_orphaned_content(storage_backend, storage_url, target_storage) + raise return file_object + @staticmethod + async def _discard_orphaned_content( + storage_backend: BaseFileStorageBackend, storage_url: str, target_storage: str + ) -> None: + try: + await storage_backend.delete_file(storage_url) + except Exception as e: # noqa: BLE001 # the metadata failure is what surfaces; a failed cleanup is only logged + verbose_proxy_logger.warning( + "Could not delete orphaned content at %s on %s after its metadata write failed: %s", + storage_url, + target_storage, + e, + ) + @staticmethod def _create_file_object_with_storage_metadata( file_content: bytes, diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 3a2ddf50143..59c0e05d97b 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -34,6 +34,7 @@ import json import logging from contextlib import ExitStack from dataclasses import dataclass +from datetime import datetime, timedelta, timezone from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -1722,6 +1723,7 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status): db_response = make_batch(id="litellm-executed-batch", status=status) db_batch_object = MagicMock() + db_batch_object.updated_at = datetime.now(timezone.utc) retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response) resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64) @@ -1734,6 +1736,41 @@ async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_ assert retrieve_harness.ensure_managed_files.call_args.kwargs["unified_batch_id"] == EXECUTED_BATCH_ID +@pytest.mark.asyncio +async def test_retrieve__executed_batch_abandoned_by_its_runner_is_served_failed(retrieve_harness, executed_runner): + runner, _ = executed_runner + failed = make_batch(id="litellm-executed-batch", status="failed") + runner.fail_abandoned = AsyncMock(return_value=failed) + db_response = make_batch(id="litellm-executed-batch", status="in_progress") + db_batch_object = MagicMock() + db_batch_object.updated_at = datetime.now(timezone.utc) - timedelta(minutes=10) + retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response) + user = UserAPIKeyAuth(api_key="sk-test", user_id="user-1") + + resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64, user=user) + + assert resp is failed + runner.fail_abandoned.assert_awaited_once_with(db_response, user) + retrieve_harness.litellm_aretrieve.assert_not_called() + retrieve_harness.router_aretrieve.assert_not_called() + retrieve_harness.ensure_managed_files.assert_called_once() + + +@pytest.mark.asyncio +async def test_retrieve__executed_batch_with_a_fresh_heartbeat_is_left_running(retrieve_harness, executed_runner): + runner, _ = executed_runner + runner.fail_abandoned = AsyncMock() + db_response = make_batch(id="litellm-executed-batch", status="in_progress") + db_batch_object = MagicMock() + db_batch_object.updated_at = datetime.now(timezone.utc) - timedelta(seconds=30) + retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response) + + resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64) + + assert resp is db_response + runner.fail_abandoned.assert_not_awaited() + + @pytest.mark.asyncio async def test_retrieve__executed_batch_without_db_row_404(retrieve_harness): with pytest.raises(ProxyException) as exc: diff --git a/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py index 96e054272a4..e5ed873a29d 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -2,6 +2,8 @@ import asyncio import json from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import Final, Literal, cast from unittest.mock import AsyncMock, MagicMock @@ -20,6 +22,7 @@ from litellm.proxy.batches_endpoints.litellm_executed_batches import ( InvalidBatchInput, LiteLLMExecutedBatchRunner, _resolve_transition, + executed_batch_runner_lost, litellm_executed_provider_for, litellm_executed_provider_of, parse_batch_input, @@ -174,10 +177,15 @@ class RealIdManagedBatchStore(FakeManagedBatchStore): class FakeManagedObjectTable: def __init__(self, objects: Mapping[str, StoredObject]) -> None: self.objects = objects + self.touches: list[tuple[str, str | None]] = [] async def find_first(self, where: Mapping[str, str]) -> StoredObject | None: return self.objects.get(where["unified_object_id"]) + async def update_many(self, where: Mapping[str, str], data: Mapping[str, str | None]) -> int: + self.touches.append((where["unified_object_id"], data["updated_by"])) + return 1 + class FakeDb: def __init__(self, objects: Mapping[str, StoredObject]) -> None: @@ -203,6 +211,9 @@ class FakeRouter: def get_model_ids(self, model_name: str) -> list[str]: return [DEPLOYMENT_ID] if model_name == BATCH_MODEL else [] + def get_model_group_info(self, model_group: str) -> None: + return None + class FakeStorageBackend: def __init__(self, contents: Mapping[str, bytes]) -> None: @@ -317,6 +328,8 @@ def make_runner( upload_error: Exception | None = None, storage_error: ValueError | None = None, store_factory: Callable[[Mapping[str, LiteLLM_ManagedFileTable]], FakeManagedBatchStore] = FakeManagedBatchStore, + general_settings: Mapping[str, object] = MappingProxyType({}), + heartbeat_seconds: float = 30.0, ) -> Harness: store = store_factory({INPUT_FILE_ID: managed_input_file()} if files is None else files) router = FakeRouter() @@ -332,7 +345,9 @@ def make_runner( prisma_client=cast("PrismaClient", prisma), managed_files=store, proxy_logging_obj=MagicMock(spec=ProxyLogging), + general_settings=general_settings, concurrency=concurrency, + heartbeat_seconds=heartbeat_seconds, storage_backend_factory=storage_factory, upload_result_file=uploads, ) @@ -413,6 +428,27 @@ def test_resolve_transition_from_cancelling(requested: BatchStatus, expected: Ba assert _resolve_transition("cancelling", requested) == expected +@pytest.mark.parametrize( + ("status", "age_seconds", "lost"), + [ + ("validating", 200, True), + ("in_progress", 200, True), + ("in_progress", 100, False), + ("finalizing", 200, True), + ("cancelling", 200, True), + ("completed", 200, False), + ("failed", 200, False), + ("cancelled", 200, False), + ("expired", 200, False), + ], +) +def test_executed_batch_runner_lost_only_for_a_stale_non_terminal_batch( + status: str, age_seconds: int, lost: bool +) -> None: + updated_at = datetime.now(timezone.utc) - timedelta(seconds=age_seconds) + assert executed_batch_runner_lost(status, updated_at) is lost + + @pytest.mark.parametrize( ("credentials", "expected"), [ @@ -686,6 +722,91 @@ async def test_create_surfaces_a_storage_backend_error_as_a_400() -> None: assert harness.store.calls == [] +CREDENTIAL_ROWS: Final = jsonl(chat_row("row-1", "hi 1"), chat_row("row-2", "hi 2", api_base="https://evil.example")) + + +async def test_create_rejects_a_row_carrying_client_side_credentials() -> None: + harness = make_runner(content=CREDENTIAL_ROWS) + with pytest.raises(ProxyException) as raised: + await harness.create() + assert raised.value.code == "400" + assert raised.value.message.startswith("Invalid batch input file: line 2") + assert "api_base" in raised.value.message + assert "allow_client_side_credentials" in raised.value.message + assert harness.store.calls == [] + assert harness.router.acompletion.await_count == 0 + + +async def test_create_forwards_row_credentials_when_the_admin_opted_in() -> None: + harness = make_runner( + content=CREDENTIAL_ROWS, general_settings=MappingProxyType({"allow_client_side_credentials": True}) + ) + _, finished = await harness.create_and_finish() + + assert finished.status == "completed" + assert finished.request_counts == BatchRequestCounts(completed=2, failed=0, total=2) + by_content = {call.kwargs["messages"][0]["content"]: call.kwargs for call in harness.router.acompletion.await_args_list} + assert by_content["hi 2"]["api_base"] == "https://evil.example" + assert "api_base" not in by_content["hi 1"] + + +async def test_running_batch_touches_its_row_until_it_finishes() -> None: + harness = make_runner(heartbeat_seconds=0.01) + + async def slow_dispatch(**_: object) -> ModelResponse: + await asyncio.sleep(0.05) + return chat_response("slow") + + harness.router.acompletion.side_effect = slow_dispatch + created, finished = await harness.create_and_finish() + + touches = harness.prisma.db.litellm_managedobjecttable.touches + assert finished.status == "completed" + assert touches + assert set(touches) == {(created.id, "user-1")} + assert [call.status for call in harness.store.calls] == ["validating", "in_progress", "finalizing", "completed"] + beats_at_finish = len(touches) + await asyncio.sleep(0.05) + assert len(touches) == beats_at_finish + + +async def test_fail_abandoned_marks_the_batch_failed_with_the_runner_lost_error() -> None: + harness = make_runner() + batch = seeded_batch(harness.store, "in_progress") + + failed = await harness.runner.fail_abandoned(batch, harness.user) + + assert failed.status == "failed" + assert failed.failed_at is not None + assert failed.errors is not None + assert [(error.message, error.code) for error in failed.errors.data or []] == [ + (litellm_executed_batches._RUNNER_LOST_MESSAGE, "runner_lost") + ] + assert harness.store.batch(batch.id).status == "failed" + assert [(call.status, call.create_if_missing) for call in harness.store.calls] == [("failed", False)] + + +async def test_running_batch_stops_and_writes_nothing_once_a_retriever_marked_it_failed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm_executed_batches, "_CANCEL_POLL_SECONDS", 0.0) + rows = jsonl(chat_row("row-1", "hi 1"), chat_row("row-2", "hi 2"), chat_row("row-3", "hi 3")) + harness = make_runner(content=rows, concurrency=1) + + def dispatch(metadata: Mapping[str, object], **_: object) -> ModelResponse: + running = harness.store.batch(str(metadata["batch_id"])) + harness.store.write(running.model_copy(update={"status": "failed"})) + return chat_response("hi 1") + + harness.router.acompletion.side_effect = dispatch + _, finished = await harness.create_and_finish() + + assert harness.router.acompletion.await_count == 1 + assert finished.status == "failed" + assert [call.status for call in harness.store.calls] == ["validating", "in_progress"] + assert harness.uploads.calls == [] + + @pytest.mark.parametrize( ("endpoint", "body", "method"), [ diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py b/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py index 067826004e2..81c5803da33 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py @@ -12,13 +12,20 @@ from litellm.proxy.utils import PrismaClient class _RecordingStorageBackend: - def __init__(self): + def __init__(self, delete_error: Exception | None = None): self.upload_calls = [] + self.delete_calls: list[str] = [] + self.delete_error = delete_error async def upload_file(self, **kwargs): self.upload_calls.append(kwargs) return "https://storage.example/blob-1" + async def delete_file(self, storage_url: str) -> None: + self.delete_calls.append(storage_url) + if self.delete_error is not None: + raise self.delete_error + class _FakeManagedFilesHook(BaseFileEndpoints): def __init__(self): @@ -45,6 +52,11 @@ class _FakeManagedFilesHook(BaseFileEndpoints): self.stored.append(kwargs) +class _FailingManagedFilesHook(_FakeManagedFilesHook): + async def store_unified_file_id(self, **kwargs): + raise RuntimeError("db down") + + class _FakeProxyLogging: def __init__(self, hook): self._hook = hook @@ -153,3 +165,25 @@ async def test_upload_hands_the_prisma_client_to_the_storage_backend_factory(mon ) assert factory_calls == [("litellm_db", prisma_client)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete_error", [None, OSError("blob locked")], ids=["delete succeeds", "delete fails"]) +async def test_upload_deletes_the_uploaded_content_when_the_metadata_write_fails( + monkeypatch: pytest.MonkeyPatch, delete_error: Exception | None +): + backend = _RecordingStorageBackend(delete_error=delete_error) + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend) + + with pytest.raises(RuntimeError, match="db down"): + await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=_FailingManagedFilesHook()), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert len(backend.upload_calls) == 1 + assert backend.delete_calls == ["https://storage.example/blob-1"]