mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(scheduler): remove a request's queue entry once it stops waiting
Requests admitted while a healthy deployment existed never left the priority queue, and neither did requests cancelled while waiting. The stale entries blocked later requests during cooldown and, with Redis, made add_request raise a TypeError on queues read back as JSON lists Both scheduling paths now share one polling helper that removes the entry in a finally block, whether the request was admitted, timed out or cancelled Related to #43059
This commit is contained in:
RachelHuangZW
parent
e64e635185
commit
876fa5510d
2 changed files with 96 additions and 101 deletions
|
|
@ -4313,57 +4313,17 @@ class Router:
|
|||
stream=False,
|
||||
**kwargs,
|
||||
):
|
||||
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
|
||||
### FLOW ITEM ###
|
||||
_request_id: Final = str(uuid.uuid4())
|
||||
item: Final = FlowItem(
|
||||
priority=priority, # 👈 SET PRIORITY FOR REQUEST
|
||||
request_id=_request_id, # 👈 SET REQUEST ID
|
||||
model_name=model, # 👈 SAME as 'Router'
|
||||
await self._wait_for_scheduler_turn(
|
||||
model=model, priority=priority, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs)
|
||||
)
|
||||
### [fin] ###
|
||||
|
||||
## ADDS REQUEST TO QUEUE ##
|
||||
await self.scheduler.add_request(request=item)
|
||||
|
||||
## POLL QUEUE
|
||||
end_time: Final = time.monotonic() + self.timeout
|
||||
curr_time = time.monotonic()
|
||||
poll_interval: Final = self.scheduler.polling_interval # poll every 3ms
|
||||
make_request = False
|
||||
|
||||
while curr_time < end_time:
|
||||
_healthy_deployments, _ = await self._async_get_healthy_deployments(
|
||||
model=model, parent_otel_span=parent_otel_span
|
||||
)
|
||||
make_request = await self.scheduler.poll( ## POLL QUEUE ## - returns 'True' if there's healthy deployments OR if request is at top of queue
|
||||
id=item.request_id,
|
||||
model_name=item.model_name,
|
||||
health_deployments=_healthy_deployments,
|
||||
)
|
||||
if make_request: ## IF TRUE -> MAKE REQUEST
|
||||
break
|
||||
else: ## ELSE -> loop till default_timeout
|
||||
await asyncio.sleep(poll_interval)
|
||||
curr_time = time.monotonic()
|
||||
|
||||
if make_request:
|
||||
try:
|
||||
_response: Final = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs)
|
||||
_response._hidden_params.setdefault("additional_headers", {})
|
||||
_response._hidden_params["additional_headers"].update({"x-litellm-request-prioritization-used": True})
|
||||
return _response
|
||||
except Exception as e:
|
||||
setattr(e, "priority", priority)
|
||||
raise e
|
||||
else:
|
||||
# Clean up the request from the scheduler queue also before raising the timeout exception
|
||||
await self.scheduler.remove_request(request_id=item.request_id, model_name=item.model_name)
|
||||
raise litellm.Timeout(
|
||||
message="Request timed out while polling queue",
|
||||
model=model,
|
||||
llm_provider="openai",
|
||||
)
|
||||
try:
|
||||
_response: Final = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs)
|
||||
_response._hidden_params.setdefault("additional_headers", {})
|
||||
_response._hidden_params["additional_headers"].update({"x-litellm-request-prioritization-used": True})
|
||||
return _response
|
||||
except Exception as e:
|
||||
setattr(e, "priority", priority)
|
||||
raise e
|
||||
|
||||
async def _schedule_factory(
|
||||
self,
|
||||
|
|
@ -4373,60 +4333,36 @@ class Router:
|
|||
args: tuple[object, ...],
|
||||
kwargs: dict[str, object],
|
||||
):
|
||||
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
|
||||
### FLOW ITEM ###
|
||||
_request_id: Final = str(uuid.uuid4())
|
||||
item: Final = FlowItem(
|
||||
priority=priority, # 👈 SET PRIORITY FOR REQUEST
|
||||
request_id=_request_id, # 👈 SET REQUEST ID
|
||||
model_name=model, # 👈 SAME as 'Router'
|
||||
await self._wait_for_scheduler_turn(
|
||||
model=model, priority=priority, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs)
|
||||
)
|
||||
### [fin] ###
|
||||
try:
|
||||
_response: Final = await original_function(*args, **kwargs)
|
||||
if isinstance(_response._hidden_params, dict):
|
||||
_response._hidden_params.setdefault("additional_headers", {})
|
||||
_response._hidden_params["additional_headers"].update({"x-litellm-request-prioritization-used": True})
|
||||
return _response
|
||||
except Exception as e:
|
||||
setattr(e, "priority", priority)
|
||||
raise e
|
||||
|
||||
## ADDS REQUEST TO QUEUE ##
|
||||
async def _wait_for_scheduler_turn(self, model: str, priority: int, parent_otel_span: Span | None) -> None:
|
||||
item: Final = FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model)
|
||||
await self.scheduler.add_request(request=item)
|
||||
|
||||
## POLL QUEUE
|
||||
end_time: Final = time.monotonic() + self.timeout
|
||||
curr_time = time.monotonic()
|
||||
poll_interval: Final = self.scheduler.polling_interval # poll every 3ms
|
||||
make_request = False
|
||||
|
||||
while curr_time < end_time:
|
||||
_healthy_deployments, _ = await self._async_get_healthy_deployments(
|
||||
model=model, parent_otel_span=parent_otel_span
|
||||
)
|
||||
make_request = await self.scheduler.poll( ## POLL QUEUE ## - returns 'True' if there's healthy deployments OR if request is at top of queue
|
||||
id=item.request_id,
|
||||
model_name=item.model_name,
|
||||
health_deployments=_healthy_deployments,
|
||||
)
|
||||
if make_request: ## IF TRUE -> MAKE REQUEST
|
||||
break
|
||||
else: ## ELSE -> loop till default_timeout
|
||||
await asyncio.sleep(poll_interval)
|
||||
curr_time = time.monotonic()
|
||||
|
||||
if make_request:
|
||||
try:
|
||||
_response: Final = await original_function(*args, **kwargs)
|
||||
if isinstance(_response._hidden_params, dict):
|
||||
_response._hidden_params.setdefault("additional_headers", {})
|
||||
_response._hidden_params["additional_headers"].update(
|
||||
{"x-litellm-request-prioritization-used": True}
|
||||
)
|
||||
return _response
|
||||
except Exception as e:
|
||||
setattr(e, "priority", priority)
|
||||
raise e
|
||||
else:
|
||||
# Clean up the request from the scheduler queue also before raising the timeout exception
|
||||
await self.scheduler.remove_request(request_id=item.request_id, model_name=item.model_name)
|
||||
raise litellm.Timeout(
|
||||
message="Request timed out while polling queue",
|
||||
model=model,
|
||||
llm_provider="openai",
|
||||
)
|
||||
try:
|
||||
end_time: Final = time.monotonic() + self.timeout
|
||||
while time.monotonic() < end_time:
|
||||
healthy_deployments, _ = await self._async_get_healthy_deployments(
|
||||
model=model, parent_otel_span=parent_otel_span
|
||||
)
|
||||
if await self.scheduler.poll(
|
||||
id=item.request_id, model_name=model, health_deployments=healthy_deployments
|
||||
):
|
||||
return
|
||||
await asyncio.sleep(self.scheduler.polling_interval)
|
||||
finally:
|
||||
await self.scheduler.remove_request(request_id=item.request_id, model_name=model)
|
||||
raise litellm.Timeout(message="Request timed out while polling queue", model=model, llm_provider="openai")
|
||||
|
||||
def _is_prompt_management_model(self, model: str) -> bool:
|
||||
model_list: Final = self.get_model_list(model_name=model)
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ 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.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
|
||||
from litellm.scheduler import FlowItem
|
||||
from litellm.types.llms.openai import ChatCompletionRequest
|
||||
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy
|
||||
|
||||
|
|
@ -17862,3 +17863,61 @@ def test_access_windows_filter_reserved_deployments_method():
|
|||
request_team_id="team-a",
|
||||
)
|
||||
] == ["reserved-deployment", "open-deployment"]
|
||||
|
||||
|
||||
def _scheduled_router(timeout: float) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "sched-model",
|
||||
"litellm_params": {"model": "openai/sched-model", "api_key": "sk-fake", "mock_response": "hi"},
|
||||
"model_info": {"id": "sched-deployment"},
|
||||
}
|
||||
],
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
async def _send_scheduled_chat(router: Router, priority: int) -> object:
|
||||
return await router.acompletion(
|
||||
model="sched-model", messages=[{"role": "user", "content": "hi"}], priority=priority
|
||||
)
|
||||
|
||||
|
||||
async def _send_scheduled_text(router: Router, priority: int) -> object:
|
||||
return await router.atext_completion(model="sched-model", prompt="hi", priority=priority)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"send", [_send_scheduled_chat, _send_scheduled_text], ids=["schedule_acompletion", "schedule_factory"]
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_admitted_prioritized_request_does_not_block_later_request_during_cooldown(
|
||||
send: Callable[[Router, int], Awaitable[object]],
|
||||
):
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
|
||||
router: Final = _scheduled_router(timeout=1)
|
||||
await send(router, 1)
|
||||
_cool_down(router, "sched-deployment")
|
||||
|
||||
with pytest.raises(RouterRateLimitError, match="cooldown"):
|
||||
await send(router, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stop_waiting", ["cancel", "timeout"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_prioritized_request_leaves_queue_when_it_stops_waiting(stop_waiting: Literal["cancel", "timeout"]):
|
||||
router: Final = _scheduled_router(timeout=0.5)
|
||||
_cool_down(router, "sched-deployment")
|
||||
await router.scheduler.add_request(FlowItem(priority=0, request_id="head-of-queue", model_name="sched-model"))
|
||||
waiting: Final = asyncio.create_task(_send_scheduled_chat(router, 5))
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(await router.scheduler.get_queue("sched-model")) == 2
|
||||
|
||||
if stop_waiting == "cancel":
|
||||
waiting.cancel()
|
||||
with pytest.raises(asyncio.CancelledError if stop_waiting == "cancel" else litellm.Timeout):
|
||||
await waiting
|
||||
|
||||
assert await router.scheduler.get_queue("sched-model") == [(0, "head-of-queue")]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue