mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
test(images): cover SDK retry budgets and router ownership
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Moe Khalil
parent
b9e3587721
commit
6b2f9d7114
3 changed files with 363 additions and 4 deletions
|
|
@ -1,6 +1,10 @@
|
|||
import json
|
||||
from typing import Final
|
||||
import threading
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
|
|
@ -11,6 +15,85 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
|||
_PROMPT: Final = "a transparent lighthouse at dawn"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("router_retries", (0, 1))
|
||||
@pytest.mark.parametrize("sdk_retries", (0, 5))
|
||||
@pytest.mark.parametrize("outcome", ("503", "timeout", "success"))
|
||||
async def test_proxy_image_retry_attempts_are_owned_only_by_router(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
router_retries: int,
|
||||
sdk_retries: int,
|
||||
outcome: Literal["503", "timeout", "success"],
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
release: Final = threading.Event()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/images/generations"
|
||||
assert request.headers["authorization"] == "Bearer synthetic-image-key"
|
||||
assert _JSON_OBJECT.validate_json(request.body) == {"model": "gpt-image-1-mini", "prompt": "a red circle"}
|
||||
if outcome == "timeout":
|
||||
return Reply(chunks=(b"{}",), gate_before_first=release)
|
||||
if outcome == "503":
|
||||
return Reply(
|
||||
status=503,
|
||||
headers=MappingProxyType({"retry-after-ms": "1"}),
|
||||
body=b'{"error":{"message":"scripted unavailable","type":"server_error"}}',
|
||||
)
|
||||
return Reply(body=b'{"created":1700000000,"data":[{"b64_json":"aW1n"}]}')
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "num_retries", None)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
with wire_server(respond) as wire:
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "image-retries",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-image-1-mini",
|
||||
"api_base": f"{wire.url}/v1",
|
||||
"api_key": "synthetic-image-key",
|
||||
"max_retries": sdk_retries,
|
||||
"timeout": 0.1,
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=router_retries,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "master_key", None)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://proxy.test"
|
||||
) as client:
|
||||
response: Final = await client.post(
|
||||
"/v1/images/generations",
|
||||
json={"model": "image-retries", "prompt": "a red circle", "max_retries": sdk_retries},
|
||||
)
|
||||
finally:
|
||||
release.set()
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
requests: Final = wire.drain()
|
||||
status: Final = {"503": 503, "timeout": 408, "success": 200}[outcome]
|
||||
attempts: Final = 1 if outcome == "success" else router_retries + 1
|
||||
assert response.status_code == status, response.text
|
||||
assert len(requests) == attempts, tuple(request.headers.get("x-stainless-retry-count") for request in requests)
|
||||
if outcome == "success":
|
||||
data: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(
|
||||
_JSON_OBJECT.validate_json(response.content)["data"]
|
||||
)
|
||||
assert data[0]["b64_json"] == "aW1n"
|
||||
|
||||
|
||||
def test_image_generation_additional_drop_params_reaches_provider_body(gateway: Gateway) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "POST"
|
||||
|
|
|
|||
|
|
@ -3,14 +3,18 @@ import base64
|
|||
import json
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from collections.abc import Generator
|
||||
from queue import SimpleQueue
|
||||
from typing import Final, Literal, TypedDict
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from openai import AsyncOpenAI
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
from openai.types.image import Image
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
|
@ -19,6 +23,211 @@ from litellm.types.utils import CallTypes, StandardLoggingPayload
|
|||
import os
|
||||
|
||||
|
||||
class ImageRetryParameters(TypedDict, total=False):
|
||||
max_retries: ReadOnly[int | None]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def image_retry_delays() -> Generator[None, None, None]:
|
||||
with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock):
|
||||
yield
|
||||
|
||||
|
||||
def _image_retry_transport(
|
||||
requests: SimpleQueue[httpx.Request], outcome: Literal["503", "timeout", "success"]
|
||||
) -> httpx.MockTransport:
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
requests.put(request)
|
||||
assert request.method == "POST"
|
||||
assert request.url.path == "/v1/images/generations"
|
||||
body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content)
|
||||
assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"}
|
||||
if outcome == "timeout":
|
||||
raise httpx.ReadTimeout("scripted image timeout", request=request)
|
||||
if outcome == "503":
|
||||
return httpx.Response(503, json={"error": {"message": "scripted unavailable", "type": "server_error"}})
|
||||
return httpx.Response(200, json={"created": 1700000000, "data": [{"b64_json": "aW1n"}]})
|
||||
|
||||
return httpx.MockTransport(respond)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("image_retry_delays")
|
||||
@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
|
||||
@pytest.mark.parametrize("outcome", ("503", "timeout"))
|
||||
@pytest.mark.parametrize(
|
||||
("retries", "attempts"),
|
||||
((0, 1), (1, 2), (None, 3), ("unset", 3), ("global", 3)),
|
||||
ids=("zero", "one", "none", "unset", "global"),
|
||||
)
|
||||
async def test_image_generation_honors_sdk_retry_budget(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
is_async: bool,
|
||||
outcome: Literal["503", "timeout"],
|
||||
retries: int | None | Literal["unset", "global"],
|
||||
attempts: int,
|
||||
) -> None:
|
||||
requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
|
||||
transport: Final = _image_retry_transport(requests, outcome)
|
||||
if retries == "global":
|
||||
monkeypatch.setattr(litellm, "max_retries", 0, raising=False)
|
||||
retry_params: Final[ImageRetryParameters] = {} if isinstance(retries, str) else {"max_retries": retries}
|
||||
error: Final = litellm.Timeout if outcome == "timeout" else litellm.ServiceUnavailableError
|
||||
if is_async:
|
||||
async with AsyncOpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
max_retries=5,
|
||||
http_client=httpx.AsyncClient(transport=transport),
|
||||
) as client:
|
||||
with pytest.raises(error):
|
||||
await litellm.aimage_generation(
|
||||
model="openai/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
**retry_params,
|
||||
)
|
||||
else:
|
||||
with OpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
max_retries=5,
|
||||
http_client=httpx.Client(transport=transport),
|
||||
) as client:
|
||||
with pytest.raises(error):
|
||||
litellm.image_generation(
|
||||
model="openai/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
**retry_params,
|
||||
)
|
||||
assert requests.qsize() == attempts
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("image_retry_delays")
|
||||
@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
|
||||
@pytest.mark.parametrize("provider", ("litellm_proxy", "hosted_vllm"))
|
||||
async def test_compatible_image_generation_honors_zero_retries(provider: str, is_async: bool) -> None:
|
||||
requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
|
||||
transport: Final = _image_retry_transport(requests, "503")
|
||||
if is_async:
|
||||
async with AsyncOpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
http_client=httpx.AsyncClient(transport=transport),
|
||||
) as client:
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
await litellm.aimage_generation(
|
||||
model=f"{provider}/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
max_retries=0,
|
||||
)
|
||||
else:
|
||||
with OpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
http_client=httpx.Client(transport=transport),
|
||||
) as client:
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
litellm.image_generation(
|
||||
model=f"{provider}/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
max_retries=0,
|
||||
)
|
||||
assert requests.qsize() == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("image_retry_delays")
|
||||
@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
|
||||
async def test_image_generation_rejects_non_integer_retries_before_sending(is_async: bool) -> None:
|
||||
requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
|
||||
transport: Final = _image_retry_transport(requests, "success")
|
||||
if is_async:
|
||||
async with AsyncOpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
http_client=httpx.AsyncClient(transport=transport),
|
||||
) as client:
|
||||
with pytest.raises(litellm.UnprocessableEntityError) as raised:
|
||||
await litellm.aimage_generation(
|
||||
model="openai/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
max_retries="invalid",
|
||||
)
|
||||
else:
|
||||
with OpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
http_client=httpx.Client(transport=transport),
|
||||
) as client:
|
||||
with pytest.raises(litellm.UnprocessableEntityError) as raised:
|
||||
litellm.image_generation(
|
||||
model="openai/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
max_retries="invalid",
|
||||
)
|
||||
assert raised.value.status_code == 422
|
||||
assert requests.qsize() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
|
||||
async def test_image_generation_returns_success_without_forwarding_retry_control(is_async: bool) -> None:
|
||||
requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
|
||||
transport: Final = _image_retry_transport(requests, "success")
|
||||
if is_async:
|
||||
async with AsyncOpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
http_client=httpx.AsyncClient(transport=transport),
|
||||
) as client:
|
||||
async_response: Final = await litellm.aimage_generation(
|
||||
model="openai/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
max_retries=0,
|
||||
)
|
||||
async_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(async_response.data)
|
||||
assert async_data[0].b64_json == "aW1n"
|
||||
else:
|
||||
with OpenAI(
|
||||
api_key="synthetic-image-key",
|
||||
base_url="https://image.test/v1",
|
||||
http_client=httpx.Client(transport=transport),
|
||||
) as client:
|
||||
sync_response: Final = litellm.image_generation(
|
||||
model="openai/gpt-image-1-mini",
|
||||
prompt="a red circle",
|
||||
api_key="synthetic-image-key",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
max_retries=0,
|
||||
)
|
||||
sync_data: Final = TypeAdapter(tuple[Image, ...]).validate_python(sync_response.data)
|
||||
assert sync_data[0].b64_json == "aW1n"
|
||||
assert requests.qsize() == 1
|
||||
|
||||
|
||||
def test_image_generation_keeps_an_internal_prefixed_kwarg_out_of_the_provider_request(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -6,7 +6,10 @@ GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params
|
|||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from unittest.mock import patch
|
||||
from queue import SimpleQueue
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
|
@ -14,6 +17,70 @@ from litellm.types.router import RetryPolicy
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", (False, True), ids=("sync", "async"))
|
||||
@pytest.mark.parametrize("source", ("deployment", "default", "request"))
|
||||
@pytest.mark.parametrize(("sdk_retries", "router_retries"), ((0, 0), (0, 1), (5, 1)))
|
||||
async def test_routed_images_do_not_multiply_sdk_and_router_retries(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
is_async: bool,
|
||||
source: Literal["deployment", "default", "request"],
|
||||
sdk_retries: int,
|
||||
router_retries: int,
|
||||
) -> None:
|
||||
requests: Final[SimpleQueue[httpx.Request]] = SimpleQueue()
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
requests.put(request)
|
||||
assert request.url.path == "/v1/images/generations"
|
||||
body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.content)
|
||||
assert body == {"model": "gpt-image-1-mini", "prompt": "a red circle"}
|
||||
return httpx.Response(
|
||||
503,
|
||||
headers={"retry-after": "0"},
|
||||
json={"error": {"message": "scripted unavailable", "type": "server_error"}},
|
||||
)
|
||||
|
||||
transport: Final = httpx.MockTransport(respond)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
with httpx.Client(transport=transport) as sync_session:
|
||||
async with httpx.AsyncClient(transport=transport) as async_session:
|
||||
monkeypatch.setattr(litellm, "client_session", sync_session)
|
||||
monkeypatch.setattr(litellm, "aclient_session", async_session)
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "num_retries", None)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "image-retries",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-image-1-mini",
|
||||
"api_key": "synthetic-image-key",
|
||||
"api_base": "https://router-image.test/v1",
|
||||
**({"max_retries": sdk_retries} if source == "deployment" else {}),
|
||||
},
|
||||
}
|
||||
],
|
||||
default_litellm_params={"max_retries": sdk_retries} if source == "default" else {},
|
||||
num_retries=router_retries,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
call_params: Final = {"max_retries": sdk_retries} if source == "request" else {}
|
||||
try:
|
||||
with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock):
|
||||
if is_async:
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
await router.aimage_generation(model="image-retries", prompt="a red circle", **call_params)
|
||||
else:
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
router.image_generation(model="image-retries", prompt="a red circle", **call_params)
|
||||
finally:
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
assert requests.qsize() == router_retries + 1
|
||||
|
||||
|
||||
class TestPerDeploymentNumRetries:
|
||||
"""Test that per-deployment num_retries in litellm_params is correctly used."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue