scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 353

fix(router): avoid coroutine leak in async pre-call check path

Browse source
This commit is contained in:
Lavya Tandel 2026-08-08 13:41:17 +05:30
parent dee667edb7
commit e1f39ab22a
4 changed files with 41 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/router.py
View file

@ -10242,7 +10242,7 @@ class Router:
return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages
raise ValueError("Either messages or input must be provided to count tokens")
def _pre_call_checks(
def _pre_call_checks_sync(
self,
model: str,
healthy_deployments: list,
@ -10419,6 +10419,29 @@ class Router:
return _returned_deployments
async def _pre_call_checks_async(
self,
model: str,
healthy_deployments: list,
messages: list[dict[str, str]] | None = None,
input: str | list | None = None,
request_kwargs: dict | None = None,
) -> list[dict]:
"""
Async wrapper around _pre_call_checks_sync.
Offloads sync work to a worker thread so async callers can await
without blocking the event loop.
"""
return await asyncio.to_thread(
self._pre_call_checks_sync,
model=model,
healthy_deployments=healthy_deployments,
messages=messages,
input=input,
request_kwargs=request_kwargs,
)
def _get_model_from_alias(self, model: str) -> str | None:
"""
Get the model from the alias.
@ -10768,7 +10791,7 @@ class Router:
)
if self.enable_pre_call_checks and (messages is not None or input is not None):
healthy_deployments = self._pre_call_checks(
healthy_deployments = await self._pre_call_checks_async(
model=model,
healthy_deployments=cast(list[dict], healthy_deployments),
messages=messages,
@ -11352,7 +11375,7 @@ class Router:
# filter pre-call checks
if self.enable_pre_call_checks and (messages is not None or input is not None):
healthy_deployments = self._pre_call_checks(
healthy_deployments = self._pre_call_checks_sync(
model=model,
healthy_deployments=healthy_deployments,
messages=messages,
@ -11510,7 +11533,7 @@ class Router:
# 5. Apply pre-call checks (if enabled)
if self.enable_pre_call_checks and (messages is not None or input is not None):
pass_through_deployments = self._pre_call_checks(
pass_through_deployments = self._pre_call_checks_sync(
model=model,
healthy_deployments=pass_through_deployments,
messages=messages,

8
tests/local_testing/test_router.py
View file

@ -634,7 +634,7 @@ def test_router_rpm_pre_call_check():
router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0) # type: ignore
try:
router._pre_call_checks(
router._pre_call_checks_sync(
model="fake-openai-endpoint",
healthy_deployments=model_list,
messages=[{"role": "user", "content": "Hey, how's it going?"}],
@ -831,7 +831,7 @@ def test_filter_invalid_params_pre_call_check():
router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0) # type: ignore
filtered_deployments = router._pre_call_checks(
filtered_deployments = router._pre_call_checks_sync(
model="gpt-3.5-turbo",
healthy_deployments=model_list,
messages=[{"role": "user", "content": "Hey, how's it going?"}],
@ -874,7 +874,7 @@ def test_router_region_pre_call_check(allowed_model_region):
router = Router(model_list=model_list, enable_pre_call_checks=True)
_healthy_deployments = router._pre_call_checks(
_healthy_deployments = router._pre_call_checks_sync(
model="gpt-3.5-turbo",
healthy_deployments=model_list,
messages=[{"role": "user", "content": "Hey!"}],
@ -1895,7 +1895,7 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider):
litellm.token_counter.side_effect = token_counter_side_effect
try:
updated_list = router._pre_call_checks(
updated_list = router._pre_call_checks_sync(
model="gpt-4",
healthy_deployments=model_list,
messages=[{"role": "user", "content": "Hey, how's it going?"}],

4
tests/router_unit_tests/test_pre_call_checks_optimization.py
View file

@ -61,7 +61,7 @@ class TestPreCallChecksOptimization:
snapshot = copy.deepcopy(deployments)
# Call the function under test
router._pre_call_checks(
router._pre_call_checks_sync(
model="gpt-5-mini",
healthy_deployments=deployments,
messages=[{"role": "user", "content": "test"}],
@ -113,7 +113,7 @@ class TestPreCallChecksOptimization:
original_large_deployment = deployments[1] # max_input_tokens=10000
# Send a long message (100 words) that exceeds 50 tokens but fits in 10000 tokens
filtered = router._pre_call_checks(
filtered = router._pre_call_checks_sync(
model="test",
healthy_deployments=deployments,
messages=[{"role": "user", "content": " ".join(["word"] * 100)}],

16
tests/test_litellm/test_router.py
View file

@ -3131,7 +3131,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch)
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}},
]
result = router._pre_call_checks(
result = router._pre_call_checks_sync(
model="m",
healthy_deployments=deployments,
messages=[{"role": "user", "content": "hi"}],
@ -3167,7 +3167,7 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}},
]
with pytest.raises(litellm.ContextWindowExceededError):
router._pre_call_checks(
router._pre_call_checks_sync(
model="m",
healthy_deployments=deployments,
messages=[{"role": "user", "content": "hi"}],
@ -3196,7 +3196,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch):
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
]
with pytest.raises(litellm.ContextWindowExceededError):
router._pre_call_checks(
router._pre_call_checks_sync(
model="m",
healthy_deployments=deployments,
input="a very long prompt that exceeds the tiny context window",
@ -3223,7 +3223,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch):
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
]
with pytest.raises(litellm.ContextWindowExceededError):
router._pre_call_checks(
router._pre_call_checks_sync(
model="m",
healthy_deployments=deployments,
input=[
@ -3263,7 +3263,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch):
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens}
)
with pytest.raises(litellm.ContextWindowExceededError):
router._pre_call_checks(
router._pre_call_checks_sync(
model="m",
healthy_deployments=deployments,
input=short_input,
@ -3324,7 +3324,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch):
deployments = [
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
]
result = router._pre_call_checks(model="m", healthy_deployments=deployments)
result = router._pre_call_checks_sync(model="m", healthy_deployments=deployments)
assert len(result) == 1
assert counted == [] # token counting skipped entirely, so no misleading error is logged
@ -7319,7 +7319,7 @@ def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(mon
"model_info": {"id": "d1"},
}
]
result = router._pre_call_checks(
result = router._pre_call_checks_sync(
model="custom-alias",
healthy_deployments=deployments,
messages=[{"role": "user", "content": "hi"}],
@ -7363,7 +7363,7 @@ def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypa
"model_info": {"id": "d1"},
}
]
result = router._pre_call_checks(
result = router._pre_call_checks_sync(
model="custom-alias",
healthy_deployments=deployments,
messages=[{"role": "user", "content": "hi"}],