From 208ef99b88690f9abb28251d87565f8887fba6ad Mon Sep 17 00:00:00 2001 From: moe-berri Date: Wed, 28 Oct 2026 10:46:47 -0700 Subject: [PATCH] fix(images): honor max_retries and prevent Router retry amplification (#48705) * test(images): cover SDK retry budgets and router ownership Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(images): forward max_retries to the OpenAI image generation client Per-call, deployment and default_litellm_params max_retries were filtered out before reaching the OpenAI-compatible image handler, which then always used 2 SDK retries. Forward the caller's value, and pin it to 0 for routed calls so the Router stays the only retry owner, matching completion() Co-authored-by: Ericz Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(images): read image max_retries through a typed kwargs view Co-authored-by: Ericz Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): assert the public invalid retry error mapping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): preserve the invalid retry status code Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): group the proxy test imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(images): discard proxy retry test routers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ericz --- litellm/images/main.py | 11 +- .../test_image_gen_drop_params_wire.py | 87 ++++++- tests/unit/images/test_main.py | 214 +++++++++++++++++- .../test_router_per_deployment_num_retries.py | 69 +++++- 4 files changed, 376 insertions(+), 5 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 9c88e737a40..3312ca417d5 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -284,6 +284,12 @@ def image_generation( extra_headers: Final = kwargs.get("extra_headers", None) headers: Final[dict] = kwargs.get("headers", None) or {} base_model: Final = kwargs.get("base_model", None) + call_kwargs: Final = cast(Mapping[str, object], kwargs) # cast-ok: image_generation's **kwargs is untyped + max_retries: Final = call_kwargs.get("max_retries") + is_router_call: Final = any( + isinstance(call_metadata := call_kwargs.get(k), Mapping) and "model_group" in call_metadata + for k in ("metadata", "litellm_metadata") + ) if extra_headers is not None: headers.update(extra_headers) model_response: ImageResponse = litellm.utils.ImageResponse() @@ -550,6 +556,7 @@ def image_generation( optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) organization: Final[str | None] = kwargs.get("organization", None) + sdk_max_retries: Final = 0 if is_router_call else max_retries model_response = openai_chat_completions.image_generation( model=model, prompt=prompt, @@ -557,7 +564,9 @@ def image_generation( api_key=api_key or dynamic_api_key, api_base=api_base, logging_obj=litellm_logging_obj, - optional_params=optional_params, + optional_params=( + optional_params if sdk_max_retries is None else {**optional_params, "max_retries": sdk_max_retries} + ), model_response=model_response, organization=organization, aimg_generation=aimg_generation, diff --git a/tests/integration/providers/test_image_gen_drop_params_wire.py b/tests/integration/providers/test_image_gen_drop_params_wire.py index 4307d20c598..f61d7d6dc17 100644 --- a/tests/integration/providers/test_image_gen_drop_params_wire.py +++ b/tests/integration/providers/test_image_gen_drop_params_wire.py @@ -1,16 +1,101 @@ import json -from typing import Final +import threading +from types import MappingProxyType +from typing import Final, Literal +import httpx import openai import pytest from integration._support.client import Gateway from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter +import litellm + _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() + router.discard() + 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" diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 38af99ff505..88d5ae1fc2e 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -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,212 @@ 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.BadRequestError) 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.BadRequestError) 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 "max retries must be an int" in raised.value.message + 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: diff --git a/tests/unit/test_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py index 99ad7c224f8..b08ae25c438 100644 --- a/tests/unit/test_router_per_deployment_num_retries.py +++ b/tests/unit/test_router_per_deployment_num_retries.py @@ -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."""