mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(batches): gate row credentials, heartbeat executed batches, clean orphaned uploads
This commit is contained in:
mateo-berri
parent
9e8b686c7a
commit
a0957edc9c
8 changed files with 391 additions and 78 deletions
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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}'.")
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue