mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
fix(rust-bridge): stamp num_retries and timeout on async native failures
The async @client wrapper sets num_retries and timeout on every exception it raises, and the Router reads num_retries off the exception to decide how many times to retry. Async native calls raised without them. Both paths now call one stamp_async_failure helper, so a native failure carries the same retry count and timeout a Python one does Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
parent
33e5d9168e
commit
e3163aa98b
4 changed files with 45 additions and 10 deletions
|
|
@ -73,7 +73,11 @@ class Feature(str, Enum):
|
|||
MESSAGES_INTERCEPTOR = "messages_interceptor"
|
||||
"""A registered MessagesInterceptor that takes over the call."""
|
||||
AGENTIC_LOOP_HOOK = "agentic_loop_hook"
|
||||
"""A logger overriding async_should_run_agentic_loop, which may run follow-up calls."""
|
||||
"""A logger overriding async_should_run_agentic_loop, which may run follow-up calls.
|
||||
|
||||
Every interception that turns a stream into a converted stream in its pre-call deployment hook
|
||||
(websearch, code interpreter, headroom) also overrides this method, so this feature keeps those
|
||||
calls on Python too."""
|
||||
TLS_POLICY = "tls_policy"
|
||||
"""An OpenSSL security level or ECDH curve that native TLS cannot apply."""
|
||||
MANAGED_RESOURCE_IDS = "managed_resource_ids"
|
||||
|
|
|
|||
|
|
@ -25,6 +25,16 @@ def call_hook(
|
|||
return hook(request, args, kwargs)
|
||||
|
||||
|
||||
async def _stamped(start: Callable[[], Awaitable[ResultT]], kwargs: Mapping[str, object]) -> ResultT:
|
||||
from litellm.utils import stamp_async_failure
|
||||
|
||||
try:
|
||||
return await start()
|
||||
except Exception as error:
|
||||
stamp_async_failure(error, kwargs)
|
||||
raise
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PublicDispatch(Generic[RequestT]):
|
||||
route: Route
|
||||
|
|
@ -100,7 +110,7 @@ class PublicDispatch(Generic[RequestT]):
|
|||
return await runtime.arun(
|
||||
self.context(native_request),
|
||||
binding=binding,
|
||||
native=lambda hook: native(hook, native_request, args, kwargs),
|
||||
native=lambda hook: _stamped(lambda: native(hook, native_request, args, kwargs), kwargs),
|
||||
python=python,
|
||||
rules=selected_rules,
|
||||
)
|
||||
|
|
@ -115,7 +125,7 @@ class PublicDispatch(Generic[RequestT]):
|
|||
return await runtime.arun(
|
||||
context,
|
||||
binding=binding,
|
||||
native=lambda hook: native(hook, request, args, kwargs),
|
||||
native=lambda hook: _stamped(lambda: native(hook, request, args, kwargs), kwargs),
|
||||
python=lambda: python(*args, **kwargs),
|
||||
rules=selected_rules,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1482,7 +1482,14 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu
|
|||
return num_retries, kwargs
|
||||
|
||||
|
||||
def _get_wrapper_timeout(kwargs: dict[str, object]) -> object:
|
||||
def stamp_async_failure(error: Exception, kwargs: Mapping[str, object]) -> None:
|
||||
num_retries: Final = kwargs.get("num_retries")
|
||||
if num_retries is not None:
|
||||
setattr(error, "num_retries", num_retries)
|
||||
setattr(error, "timeout", _get_wrapper_timeout(kwargs=kwargs))
|
||||
|
||||
|
||||
def _get_wrapper_timeout(kwargs: Mapping[str, object]) -> object:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
if kwargs.get("stream") is not True:
|
||||
|
|
@ -2350,12 +2357,7 @@ def client(original_function):
|
|||
else:
|
||||
return result
|
||||
|
||||
deployment_num_retries: Final = kwargs.get("num_retries")
|
||||
if deployment_num_retries is not None:
|
||||
setattr(e, "num_retries", deployment_num_retries)
|
||||
|
||||
timeout: Final = _get_wrapper_timeout(kwargs=kwargs)
|
||||
setattr(e, "timeout", timeout)
|
||||
stamp_async_failure(e, kwargs)
|
||||
raise e
|
||||
|
||||
finally:
|
||||
|
|
|
|||
|
|
@ -285,3 +285,22 @@ async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Ent
|
|||
assert inspect.iscoroutine(pending)
|
||||
with pytest.raises(RuntimeError, match="is unavailable"):
|
||||
await pending
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id)
|
||||
async def test_async_native_failures_carry_the_retry_count_and_timeout_the_router_reads(entry: Entry) -> None:
|
||||
failure: Final = litellm.RateLimitError(message="limited", llm_provider="openai", model="gpt-4o")
|
||||
|
||||
async def failing(_: NativeCall) -> object:
|
||||
raise failure
|
||||
|
||||
entry.async_binding.override(failing)
|
||||
try:
|
||||
with scenario_rules(entry, SCENARIOS["required"]), pytest.raises(litellm.RateLimitError) as caught:
|
||||
await entry.asynchronous(*entry.args(), **entry.kwargs(num_retries=3, timeout=7.0))
|
||||
finally:
|
||||
entry.async_binding.reset()
|
||||
|
||||
assert caught.value is failure
|
||||
assert getattr(caught.value, "num_retries", None) == 3
|
||||
assert getattr(caught.value, "timeout", None) == 7.0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue