scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

Merge pull request #41555 from BerriAI/litellm_max_parallel_requests_queue_size

Browse source 
feat(router): reject with 429 when a deployment's max_parallel_requests slots are all in use
This commit is contained in:
Yassin Kortam 2026-09-18 11:18:08 -07:00 • committed by GitHub
commit 660f3dfdd3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 567 additions and 352 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

446
litellm/router.py
View file

@ -155,7 +155,7 @@ from litellm.router_utils.batch_utils import (
replace_model_in_jsonl,
should_replace_model_in_jsonl,
)
from litellm.router_utils.client_initalization_utils import InitalizeCachedClient
from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
is_clientside_credential,
@ -3642,24 +3642,22 @@ class Router:
input_kwargs.pop("silent_model", None)
input_kwargs.pop("include_fallback_errors", None)
_response: Final = litellm.acompletion(**input_kwargs)
logging_obj: Final[LiteLLMLogging | None] = kwargs.get("litellm_logging_obj", None)
rpm_semaphore: Final = self._get_client(
max_parallel_requests_limit: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
async with contextlib.AsyncExitStack() as deployment_slot:
if isinstance(rpm_semaphore, asyncio.Semaphore):
await deployment_slot.enter_async_context(rpm_semaphore)
if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit):
deployment_slot.enter_context(max_parallel_requests_limit)
await self.async_routing_strategy_pre_call_checks(
deployment=deployment,
logging_obj=logging_obj,
parent_otel_span=parent_otel_span,
)
response = await _response
response = await litellm.acompletion(**input_kwargs)
## CHECK CONTENT FILTER ERROR ##
if isinstance(response, ModelResponse):
@ -4586,38 +4584,16 @@ class Router:
)
self.total_calls[model_name] += 1
response = litellm.aimage_generation(
**{
**data,
"prompt": prompt,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
### CONCURRENCY-SAFE RPM CHECKS ###
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.aimage_generation(
**{
**data,
"prompt": prompt,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -4691,38 +4667,16 @@ class Router:
)
self.total_calls[model_name] += 1
response = litellm.atranscription(
**{
**data,
"file": file,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
### CONCURRENCY-SAFE RPM CHECKS ###
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.atranscription(
**{
**data,
"file": file,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -4806,38 +4760,16 @@ class Router:
)
self.total_calls[model_name] += 1
response = litellm.aspeech(
**{
**data,
"input": input,
"voice": data.get("voice") if voice is None else voice,
"client": model_client,
**kwargs,
}
)
### CONCURRENCY-SAFE RPM CHECKS ###
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.aspeech(
**{
**data,
"input": input,
"voice": data.get("voice") if voice is None else voice,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -5002,37 +4934,16 @@ class Router:
)
self.total_calls[model_name] += 1
response = litellm.atext_completion(
**{
**data,
"prompt": prompt,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.atext_completion(
**{
**data,
"prompt": prompt,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -5093,37 +5004,16 @@ class Router:
)
self.total_calls[model_name] += 1
response = litellm.aadapter_completion(
**{
**data,
"adapter_id": adapter_id,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.aadapter_completion(
**{
**data,
"adapter_id": adapter_id,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -5353,29 +5243,8 @@ class Router:
if custom_llm_provider is not None:
response_kwargs["custom_llm_provider"] = custom_llm_provider
response = original_generic_function(**response_kwargs)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await original_generic_function(**response_kwargs)
if self._should_raise_anthropic_refusal_error(
model=model,
@ -5983,38 +5852,16 @@ class Router:
)
self.total_calls[model_name] += 1
response = litellm.aembedding(
**{
**data,
"input": input,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
### CONCURRENCY-SAFE RPM CHECKS ###
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.aembedding(
**{
**data,
"input": input,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -6123,37 +5970,18 @@ class Router:
"gcs_bucket_name" in data
): # TODO: Remove this once we have a better way to handle GCS bucket name: Problem is that we need to pass the gcs_bucket_name to the router for the create_file call but it doesn't show up there
kwargs_copy.setdefault("litellm_metadata", {})["gcs_bucket_name"] = data["gcs_bucket_name"]
response = litellm.acreate_file(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs_copy,
}
)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs_copy,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(
deployment=deployment, kwargs=kwargs_copy, parent_otel_span=parent_otel_span
):
response = await litellm.acreate_file(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs_copy,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acreate_file(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -6243,33 +6071,16 @@ class Router:
)
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
response = avector_store_create_sdk(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await avector_store_create_sdk(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -6355,37 +6166,16 @@ class Router:
)
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
response = litellm.acreate_batch(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.acreate_batch(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acreate_batch(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -6576,37 +6366,16 @@ class Router:
)
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
response = litellm.acancel_batch(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
rpm_semaphore: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
)
response = await response
else:
await self.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=parent_otel_span
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await litellm.acancel_batch(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
"client": model_client,
**kwargs,
}
)
response = await response
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acancel_batch(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@ -8741,6 +8510,23 @@ class Router:
)
raise e
@contextlib.asynccontextmanager
async def _deployment_slot(
self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None
) -> AsyncGenerator[None, None]:
"""Holds the deployment's max_parallel_requests slot, if it has one, around the provider call. Routing
strategy pre-call checks run inside the slot so their rpm accounting stays concurrency-safe."""
max_parallel_requests_limit: Final = self._get_client(
deployment=deployment,
kwargs=kwargs,
client_type="max_parallel_requests",
)
async with contextlib.AsyncExitStack() as slot:
if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit):
slot.enter_context(max_parallel_requests_limit)
await self.async_routing_strategy_pre_call_checks(deployment=deployment, parent_otel_span=parent_otel_span)
yield
async def async_callback_filter_deployments(
self,
model: str,

49
litellm/router_utils/client_initalization_utils.py
View file

@ -1,6 +1,8 @@
import asyncio
from types import TracebackType
from typing import TYPE_CHECKING, Any, Final
from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType
from litellm.types.router import RouterErrors
from litellm.utils import calculate_max_parallel_requests
if TYPE_CHECKING:
@ -11,6 +13,43 @@ else:
LitellmRouter = Any
class MaxParallelRequestsLimit:
"""A deployment's max_parallel_requests slots. A caller arriving while every slot is in use gets a 429 instead
of waiting for one to free up."""
def __init__(self, max_parallel_requests: int, model_id: str, model_group: str) -> None:
self.max_parallel_requests: Final = max_parallel_requests
self.model_id: Final = model_id
self.model_group: Final = model_group
self.in_flight = 0
def __enter__(self) -> None:
self.acquire()
def __exit__(
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
) -> None:
self.release()
def acquire(self) -> None:
if self.in_flight >= self.max_parallel_requests:
raise RateLimitError(
message=(
f"{RouterErrors.max_parallel_requests_exceeded.value} Deployment model_group={self.model_group}, "
f"id={self.model_id} already has max_parallel_requests={self.max_parallel_requests} requests in "
"flight. Raise max_parallel_requests (or the rpm/tpm it is derived from) for this deployment"
),
llm_provider="",
model=self.model_group,
category=RateLimitErrorCategory.LITELLM_RATE_LIMIT,
rate_limit_type=RateLimitType.CONCURRENT_REQUESTS,
)
self.in_flight += 1
def release(self) -> None:
self.in_flight -= 1
class InitalizeCachedClient:
@staticmethod
def set_max_parallel_requests_client(litellm_router_instance: LitellmRouter, model: dict):
@ -26,10 +65,14 @@ class InitalizeCachedClient:
default_max_parallel_requests=litellm_router_instance.default_max_parallel_requests,
)
if calculated_max_parallel_requests:
semaphore: Final = asyncio.Semaphore(calculated_max_parallel_requests)
limit: Final = MaxParallelRequestsLimit(
max_parallel_requests=calculated_max_parallel_requests,
model_id=model_id,
model_group=model.get("model_name", ""),
)
cache_key: Final = f"{model_id}_max_parallel_requests_client"
litellm_router_instance.cache.set_cache(
key=cache_key,
value=semaphore,
value=limit,
local_only=True,
)

1
litellm/types/router.py
View file

@ -653,6 +653,7 @@ class RouterErrors(enum.Enum):
"""
user_defined_ratelimit_error = "Deployment over user-defined ratelimit."
max_parallel_requests_exceeded = "Deployment has all max_parallel_requests slots in use."
no_deployments_available = "No deployments available for selected model"
all_deployments_in_cooldown = "All deployments for selected model are in cooldown"
no_deployments_with_tag_routing = "Not allowed to access model due to tags configuration"

13
tests/local_testing/test_router_max_parallel_requests.py
View file

@ -11,6 +11,7 @@ import pytest
from typing import Optional
import litellm
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
from litellm.utils import calculate_max_parallel_requests
"""
@ -93,26 +94,26 @@ def test_setting_mpr_limits_per_model(
default_max_parallel_requests=default_max_parallel_requests,
)
mpr_client: Optional[asyncio.Semaphore] = router._get_client(
mpr_client: Optional[MaxParallelRequestsLimit] = router._get_client(
deployment=deployment,
kwargs={},
client_type="max_parallel_requests",
)
if max_parallel_requests is not None:
assert max_parallel_requests == mpr_client._value
assert max_parallel_requests == mpr_client.max_parallel_requests
elif rpm is not None:
assert rpm == mpr_client._value
assert rpm == mpr_client.max_parallel_requests
elif tpm is not None:
calculated_rpm = int(tpm / 1000 * 6)
if calculated_rpm == 0:
calculated_rpm = 1
print(
f"test calculated_rpm: {calculated_rpm}, calculated_max_parallel_requests={mpr_client._value}"
f"test calculated_rpm: {calculated_rpm}, calculated_max_parallel_requests={mpr_client.max_parallel_requests}"
)
assert calculated_rpm == mpr_client._value
assert calculated_rpm == mpr_client.max_parallel_requests
elif default_max_parallel_requests is not None:
assert mpr_client._value == default_max_parallel_requests
assert mpr_client.max_parallel_requests == default_max_parallel_requests
else:
assert mpr_client is None

125
tests/test_litellm/router_utils/test_client_initalization_utils.py Normal file
View file

@ -0,0 +1,125 @@
import asyncio
from typing import Final
import pytest
import litellm
from litellm import Router
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
def _limit(max_parallel_requests: int = 1) -> MaxParallelRequestsLimit:
return MaxParallelRequestsLimit(
max_parallel_requests=max_parallel_requests, model_id="deployment-1", model_group="gpt-5.6"
)
async def _hold(limit: MaxParallelRequestsLimit, release: asyncio.Event) -> str:
with limit:
await release.wait()
return "ok"
def _expect_rejection(limit: MaxParallelRequestsLimit) -> litellm.RateLimitError:
with pytest.raises(litellm.RateLimitError) as excinfo:
limit.acquire()
return excinfo.value
@pytest.mark.asyncio
async def test_request_arriving_while_every_slot_is_in_use_gets_429_without_waiting():
limit: Final = _limit(max_parallel_requests=2)
release: Final = asyncio.Event()
holders: Final = [asyncio.create_task(_hold(limit, release)) for _ in range(2)]
await asyncio.sleep(0)
assert limit.in_flight == 2
rejection: Final = _expect_rejection(limit)
assert rejection.status_code == 429
assert "deployment-1" in rejection.message
assert "gpt-5.6" in rejection.message
assert "max_parallel_requests=2" in rejection.message
assert limit.in_flight == 2
release.set()
assert await asyncio.wait_for(asyncio.gather(*holders), timeout=2) == ["ok", "ok"]
assert limit.in_flight == 0
with limit:
assert limit.in_flight == 1
assert limit.in_flight == 0
@pytest.mark.asyncio
async def test_burst_over_the_cap_admits_exactly_max_parallel_requests_and_rejects_the_rest():
limit: Final = _limit(max_parallel_requests=3)
release: Final = asyncio.Event()
async def attempt() -> str:
try:
return await _hold(limit, release)
except litellm.RateLimitError as e:
return f"rejected:{e.status_code}"
callers: Final = [asyncio.create_task(attempt()) for _ in range(10)]
await asyncio.sleep(0)
assert limit.in_flight == 3
release.set()
outcomes: Final = await asyncio.wait_for(asyncio.gather(*callers), timeout=2)
assert outcomes.count("ok") == 3
assert outcomes.count("rejected:429") == 7
assert limit.in_flight == 0
def test_slot_is_released_when_the_held_call_raises():
limit: Final = _limit()
with pytest.raises(RuntimeError):
with limit:
raise RuntimeError("provider blew up")
assert limit.in_flight == 0
with limit:
assert limit.in_flight == 1
def _router_limit(router: Router, model_name: str) -> MaxParallelRequestsLimit:
deployment: Final = router.get_deployment_by_model_group_name(model_group_name=model_name)
assert deployment is not None
client: Final = router._get_client(
deployment=deployment.model_dump(), kwargs={}, client_type="max_parallel_requests"
)
assert isinstance(client, MaxParallelRequestsLimit)
return client
@pytest.mark.parametrize(
("litellm_params", "expected_cap"),
[
({"max_parallel_requests": 2, "rpm": 7, "tpm": 100_000}, 2),
({"rpm": 7, "tpm": 100_000}, 7),
({"tpm": 100_000}, 600),
({"tpm": 100}, 1),
],
)
@pytest.mark.asyncio
async def test_router_deployment_rejects_past_its_derived_cap(litellm_params: dict[str, int], expected_cap: int):
router: Final = Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", **litellm_params}}]
)
limit: Final = _router_limit(router, "gpt-5.6")
assert limit.max_parallel_requests == expected_cap
release: Final = asyncio.Event()
holders: Final = [asyncio.create_task(_hold(limit, release)) for _ in range(expected_cap)]
await asyncio.sleep(0)
assert limit.in_flight == expected_cap
assert f"max_parallel_requests={expected_cap}" in _expect_rejection(limit).message
release.set()
assert await asyncio.wait_for(asyncio.gather(*holders), timeout=2) == ["ok"] * expected_cap
def test_router_without_any_concurrency_setting_has_no_limit():
router: Final = Router(model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6"}}])
deployment: Final = router.get_deployment_by_model_group_name(model_group_name="gpt-5.6")
assert deployment is not None
assert (
router._get_client(deployment=deployment.model_dump(), kwargs={}, client_type="max_parallel_requests") is None
)

257
tests/test_litellm/test_router.py
View file

@ -1,11 +1,13 @@
import asyncio
import copy
import functools
import gc
import json
import logging
import os
import sys
import threading
import warnings
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime, timedelta
from types import SimpleNamespace
@ -45,6 +47,8 @@ from litellm.router import (
_is_retriable_anthropic_status,
)
from litellm.router_strategy import simple_shuffle
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments
from litellm.types.llms.openai import ChatCompletionRequest
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy
@ -1517,7 +1521,9 @@ async def test_router_ageneric_api_call_with_fallbacks_helper():
},
}
mock_semaphore = asyncio.Semaphore(1)
mock_semaphore = MaxParallelRequestsLimit(
max_parallel_requests=1, model_id="deployment-1", model_group="gpt-3.5-turbo"
)
with patch.object(
router, "_update_kwargs_with_deployment"
@ -15962,7 +15968,7 @@ def _max_parallel_router(max_parallel_requests: int) -> Router:
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_router_max_parallel_requests_bounds_in_flight_upstream_calls(
async def test_router_max_parallel_requests_admits_the_cap_and_rejects_the_rest_with_429(
monkeypatch: pytest.MonkeyPatch, stream: bool
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
@ -15988,24 +15994,33 @@ async def test_router_max_parallel_requests_bounds_in_flight_upstream_calls(
},
)
async def one_call() -> None:
response = await router.acompletion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=stream
)
async def one_call() -> str:
try:
response = await router.acompletion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=stream
)
except litellm.RateLimitError as e:
return f"rejected:{e.status_code}"
if stream:
async for _ in response:
pass
return "ok"
with respx.mock(assert_all_called=True) as respx_mock:
respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream)
await asyncio.wait_for(asyncio.gather(*(one_call() for _ in range(10))), timeout=10)
route: Final = respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream)
outcomes: Final = await asyncio.wait_for(asyncio.gather(*(one_call() for _ in range(10))), timeout=10)
assert tracker.peak <= 2
assert outcomes.count("ok") == 2
assert outcomes.count("rejected:429") == 8
assert route.call_count == 2
assert tracker.peak == 2
assert tracker.current == 0
@pytest.mark.asyncio
async def test_router_max_parallel_requests_slot_released_when_stream_closed_early(monkeypatch: pytest.MonkeyPatch):
async def test_router_max_parallel_requests_slot_held_until_stream_closed_then_released(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
tracker: Final = _InFlightTracker()
router: Final = _max_parallel_router(max_parallel_requests=1)
@ -16028,16 +16043,232 @@ async def test_router_max_parallel_requests_slot_released_when_stream_closed_ear
async for _ in second:
pass
second_task: Final = asyncio.create_task(second_call())
await asyncio.sleep(0.05)
assert tracker.current == 1
with pytest.raises(litellm.RateLimitError) as while_streaming:
await second_call()
assert while_streaming.value.status_code == 429
await first.aclose()
await asyncio.wait_for(second_task, timeout=2)
await asyncio.wait_for(second_call(), timeout=2)
assert tracker.peak == 1
assert tracker.current == 0
@pytest.mark.asyncio
async def test_router_max_parallel_requests_overflow_is_429_without_cooldown_or_provider_call(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router: Final = Router(
model_list=[
{
"model_name": "gpt-5.6",
"litellm_params": {
"model": "openai/gpt-5.6",
"api_key": "sk-fake",
"api_base": "https://max-parallel.local/v1",
"max_parallel_requests": 1,
},
"model_info": {"id": "capped-deployment"},
},
{
"model_name": "gpt-5.6",
"litellm_params": {
"model": "openai/gpt-5.6",
"api_key": "sk-fake",
"api_base": "https://max-parallel-sibling.local/v1",
},
"model_info": {"id": "sibling-deployment"},
},
],
num_retries=0,
)
async def upstream(request: httpx.Request) -> httpx.Response:
await asyncio.sleep(0.2)
return httpx.Response(
200,
json={
"id": "c",
"object": "chat.completion",
"created": 1,
"model": "gpt-5.6",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}],
},
)
with respx.mock(assert_all_called=False) as respx_mock:
route: Final = respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream)
sibling_route: Final = respx_mock.post("https://max-parallel-sibling.local/v1/chat/completions").mock(
side_effect=upstream
)
results: Final = await asyncio.wait_for(
asyncio.gather(
*(
router.acompletion(model="capped-deployment", messages=[{"role": "user", "content": "hi"}])
for _ in range(3)
),
return_exceptions=True,
),
timeout=10,
)
rejected: Final = [r for r in results if isinstance(r, BaseException)]
assert len(rejected) == 2 and len(results) == 3
assert all(isinstance(r, litellm.RateLimitError) and r.status_code == 429 for r in rejected)
assert all("capped-deployment" in r.message and "max_parallel_requests=1" in r.message for r in rejected)
assert route.call_count == 1
assert sibling_route.call_count == 0
assert await _async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) == []
@pytest.mark.asyncio
async def test_router_embedding_path_rejects_past_max_parallel_requests_without_orphan_coroutines(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router: Final = Router(
model_list=[
{
"model_name": "embed",
"litellm_params": {
"model": "openai/text-embedding-3-small",
"api_key": "sk-fake",
"api_base": "https://max-parallel-embed.local/v1",
"max_parallel_requests": 1,
},
"model_info": {"id": "embed-capped-deployment"},
}
],
num_retries=0,
)
async def upstream(request: httpx.Request) -> httpx.Response:
await asyncio.sleep(0.2)
return httpx.Response(
200,
json={
"object": "list",
"data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
"model": "text-embedding-3-small",
"usage": {"prompt_tokens": 1, "total_tokens": 1},
},
)
with respx.mock() as respx_mock, warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
route: Final = respx_mock.post("https://max-parallel-embed.local/v1/embeddings").mock(side_effect=upstream)
results: Final = await asyncio.wait_for(
asyncio.gather(
*(router.aembedding(model="embed", input=["hi"]) for _ in range(3)),
return_exceptions=True,
),
timeout=10,
)
gc.collect()
rejected: Final = [r for r in results if isinstance(r, BaseException)]
assert len(rejected) == 2 and len(results) == 3
assert all(isinstance(r, litellm.RateLimitError) and r.status_code == 429 for r in rejected)
assert all("embed-capped-deployment" in r.message for r in rejected)
assert route.call_count == 1
assert [str(w.message) for w in caught if "never awaited" in str(w.message)] == []
@pytest.mark.asyncio
async def test_router_max_parallel_requests_overflow_takes_the_ordinary_429_fallback_path(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router: Final = Router(
model_list=[
{
"model_name": "gpt-5.6",
"litellm_params": {
"model": "openai/gpt-5.6",
"api_key": "sk-fake",
"api_base": "https://max-parallel-primary.local/v1",
"max_parallel_requests": 1,
},
"model_info": {"id": "capped-primary-deployment"},
},
{
"model_name": "gpt-5.6-fallback",
"litellm_params": {
"model": "openai/gpt-5.6",
"api_key": "sk-fake",
"api_base": "https://max-parallel-fallback.local/v1",
},
"model_info": {"id": "fallback-deployment"},
},
],
fallbacks=[{"gpt-5.6": ["gpt-5.6-fallback"]}],
num_retries=0,
)
async def upstream(request: httpx.Request) -> httpx.Response:
await asyncio.sleep(0.2)
return httpx.Response(
200,
json={
"id": "c",
"object": "chat.completion",
"created": 1,
"model": "gpt-5.6",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}],
},
)
with respx.mock() as respx_mock:
primary: Final = respx_mock.post("https://max-parallel-primary.local/v1/chat/completions").mock(
side_effect=upstream
)
fallback: Final = respx_mock.post("https://max-parallel-fallback.local/v1/chat/completions").mock(
side_effect=upstream
)
results: Final = await asyncio.wait_for(
asyncio.gather(
*(router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) for _ in range(3))
),
timeout=10,
)
assert len(results) == 3
assert primary.call_count == 1
assert fallback.call_count == 2
assert await _async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) == []
@pytest.mark.asyncio
async def test_router_deployment_slot_rejects_while_held_and_frees_slot_on_exit():
router: Final = Router(
model_list=[
{
"model_name": "gpt-5.6",
"litellm_params": {
"model": "openai/gpt-5.6",
"api_key": "sk-fake",
"max_parallel_requests": 1,
},
"model_info": {"id": "slot-deployment"},
}
]
)
deployment: Final = router.get_deployment(model_id="slot-deployment")
assert deployment is not None
kwargs: Final = {"model": "gpt-5.6"}
async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None):
with pytest.raises(litellm.RateLimitError) as overflow:
async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None):
pass
assert overflow.value.status_code == 429
assert "slot-deployment" in overflow.value.message
async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None):
pass
@pytest.mark.asyncio
async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch):
from litellm import Router

28
tests/test_litellm/test_utils.py
View file

@ -61,6 +61,7 @@ from litellm.utils import (
_snapshot_exception_for_hook,
async_post_call_failure_deployment_hook,
async_post_call_success_deployment_hook,
calculate_max_parallel_requests,
client,
get_non_default_completion_params,
get_optional_params_image_gen,
@ -5692,3 +5693,30 @@ def test_get_model_info_gemini(monkeypatch):
assert info.get("rpm") is not None, f"{model} does not have rpm"
@pytest.mark.parametrize(
("max_parallel_requests", "rpm", "tpm", "default_max_parallel_requests", "expected"),
[
(3, 100, 100_000, 7, 3),
(None, 100, 100_000, 7, 100),
(None, None, 100_000, 7, 600),
(None, None, 50, 7, 1),
(None, None, None, 7, 7),
(None, None, None, None, None),
],
)
def test_calculate_max_parallel_requests_precedence(
max_parallel_requests: int | None,
rpm: int | None,
tpm: int | None,
default_max_parallel_requests: int | None,
expected: int | None,
) -> None:
assert (
calculate_max_parallel_requests(
max_parallel_requests=max_parallel_requests,
rpm=rpm,
tpm=tpm,
default_max_parallel_requests=default_max_parallel_requests,
)
== expected
)