From 79287a7584d89e2511959f3e9ab38425b4673c10 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 31 May 2024 18:51:13 -0700 Subject: [PATCH 1/8] feat(scheduler.py): add request prioritization scheduler allow user to set priority for a request --- .../out/{404.html => 404/index.html} | 0 .../{model_hub.html => model_hub/index.html} | 0 litellm/proxy/_super_secret_config.yaml | 12 +- litellm/proxy/proxy_server.py | 180 ++++++------------ litellm/proxy/queue/scheduler.py | 129 +++++++++++++ litellm/tests/test_scheduler.py | 164 ++++++++++++++++ litellm/types/router.py | 2 + tests/test_openai_endpoints.py | 30 +++ 8 files changed, 394 insertions(+), 123 deletions(-) rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) create mode 100644 litellm/proxy/queue/scheduler.py create mode 100644 litellm/tests/test_scheduler.py diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_super_secret_config.yaml b/litellm/proxy/_super_secret_config.yaml index 00efca52255..97408262c25 100644 --- a/litellm/proxy/_super_secret_config.yaml +++ b/litellm/proxy/_super_secret_config.yaml @@ -5,12 +5,12 @@ model_list: model: openai/my-fake-model rpm: 800 model_name: gpt-3.5-turbo-fake-model -- litellm_params: - api_base: https://my-endpoint-europe-berri-992.openai.azure.com/ - api_key: os.environ/AZURE_EUROPE_API_KEY - model: azure/gpt-35-turbo - rpm: 10 - model_name: gpt-3.5-turbo-fake-model +# - litellm_params: +# api_base: https://my-endpoint-europe-berri-992.openai.azure.com/ +# api_key: os.environ/AZURE_EUROPE_API_KEY +# model: azure/gpt-35-turbo +# rpm: 10 +# model_name: gpt-3.5-turbo-fake-model - litellm_params: api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ api_key: os.environ/AZURE_API_KEY diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b6e27cba60e..3a0efd653f3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -141,6 +141,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.exceptions import RejectedRequestError from litellm.integrations.slack_alerting import SlackAlertingArgs, SlackAlerting +from litellm.proxy.queue.scheduler import Scheduler, FlowItem, DefaultPriorities try: from litellm._version import version @@ -395,6 +396,8 @@ proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) async_result = None celery_app_conn = None celery_fn = None # Redis Queue for handling requests +### SIMPLE QUEUE ### +scheduler = Scheduler() ### DB WRITER ### db_writer_client: Optional[HTTPHandler] = None ### logger ### @@ -3655,7 +3658,7 @@ def on_backoff(details): @router.on_event("startup") async def startup_event(): - global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db + global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, scheduler import json ### LOAD MASTER KEY ### @@ -3691,6 +3694,10 @@ async def startup_event(): ## Error Tracking ## error_tracking() + ## Priority Workload Scheduler ## + if llm_router is not None: + scheduler.update_variables(llm_router=llm_router) + ## UPDATE SLACK ALERTING ## proxy_logging_obj.slack_alerting_instance.update_values(llm_router=llm_router) @@ -11219,118 +11226,7 @@ async def alerting_settings( return return_val -# @router.post( -# "/alerting/update", -# description="Update the slack alerting settings. Persist value in db.", -# tags=["alerting"], -# dependencies=[Depends(user_api_key_auth)], -# include_in_schema=False, -# ) -# async def alerting_update( -# data: SlackAlertingArgs, -# user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -# ): -# """Allows updating slack alerting values. Used by UI.""" -# global prisma_client -# if prisma_client is None: -# raise HTTPException( -# status_code=400, -# detail={"error": CommonProxyErrors.db_not_connected_error.value}, -# ) - -# if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: -# raise HTTPException( -# status_code=400, -# detail={"error": CommonProxyErrors.not_allowed_access.value}, -# ) - -# ## get general settings from db -# db_general_settings = await prisma_client.db.litellm_config.find_first( -# where={"param_name": "general_settings"} -# ) -# ### update value - -# alerting_args_dict = {} -# if db_general_settings is None or db_general_settings.param_value is None: -# general_settings = {} -# alerting_args_dict = {} -# else: -# general_settings = dict(db_general_settings.param_value) -# _alerting_args_dict = general_settings.get("alerting_args", None) -# if _alerting_args_dict is not None and isinstance(_alerting_args_dict, dict): -# alerting_args_dict = _alerting_args_dict - - -# alerting_args_dict = data.model - -# response = await prisma_client.db.litellm_config.upsert( -# where={"param_name": "general_settings"}, -# data={ -# "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore -# "update": {"param_value": json.dumps(general_settings)}, # type: ignore -# }, -# ) - -# return response - - #### EXPERIMENTAL QUEUING #### -async def _litellm_chat_completions_worker(data, user_api_key_dict): - """ - worker to make litellm completions calls - """ - while True: - try: - ### CALL HOOKS ### - modify incoming data before calling the model - data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion" - ) - - verbose_proxy_logger.debug("_litellm_chat_completions_worker started") - ### ROUTE THE REQUEST ### - router_model_names = ( - [m["model_name"] for m in llm_model_list] - if llm_model_list is not None - else [] - ) - if ( - llm_router is not None and data["model"] in router_model_names - ): # model in router model list - response = await llm_router.acompletion(**data) - elif ( - llm_router is not None and data["model"] in llm_router.deployment_names - ): # model in router deployments, calling a specific deployment on the router - response = await llm_router.acompletion( - **data, specific_deployment=True - ) - elif ( - llm_router is not None - and llm_router.model_group_alias is not None - and data["model"] in llm_router.model_group_alias - ): # model set in model_group_alias - response = await llm_router.acompletion(**data) - else: # router is not set - response = await litellm.acompletion(**data) - - verbose_proxy_logger.debug("final response: {response}") - return response - except HTTPException as e: - verbose_proxy_logger.debug( - f"EXCEPTION RAISED IN _litellm_chat_completions_worker - {e.status_code}; {e.detail}" - ) - if ( - e.status_code == 429 - and "Max parallel request limit reached" in e.detail - ): - verbose_proxy_logger.debug("Max parallel request limit reached!") - timeout = litellm._calculate_retry_after( - remaining_retries=3, max_retries=3, min_timeout=1 - ) - await asyncio.sleep(timeout) - else: - raise e - - @router.post( "/queue/chat/completions", tags=["experimental"], @@ -11338,6 +11234,7 @@ async def _litellm_chat_completions_worker(data, user_api_key_dict): ) async def async_queue_request( request: Request, + fastapi_response: Response, model: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): @@ -11403,12 +11300,47 @@ async def async_queue_request( if user_api_base: data["api_base"] = user_api_base - response = await asyncio.wait_for( - _litellm_chat_completions_worker( - data=data, user_api_key_dict=user_api_key_dict - ), - timeout=litellm.request_timeout, + ## FLOW ITEM ## + request_id = str(uuid.uuid4()) + flow_item = FlowItem( + priority=data.pop("priority", DefaultPriorities.Medium.value), + request_id=request_id, + model_group=data["model"], ) + # [TODO] only allow premium users to set non default priorities + + ## ADD REQUEST TO QUEUE + response = await scheduler.add_request(request=flow_item) + + if llm_router is None: + raise HTTPException( + status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} + ) + ## POLL QUEUE + default_timeout = llm_router.timeout + end_time = time.time() + default_timeout + poll_interval = 0.03 # poll every 3ms + curr_time = time.time() + + make_request = False + + if llm_router is None: + raise HTTPException( + status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} + ) + + while curr_time < end_time: + make_request = await scheduler.poll( + id=request_id, model_group=data["model"] + ) + if make_request: ## IF TRUE -> MAKE REQUEST + break + else: ## ELSE -> loop till default_timeout + await asyncio.sleep(poll_interval) + curr_time = time.time() + + if make_request: + response = await llm_router.acompletion(**data) if ( "stream" in data and data["stream"] == True @@ -11422,6 +11354,7 @@ async def async_queue_request( media_type="text/event-stream", ) + fastapi_response.headers.update({"x-litellm-priority": str(flow_item.priority)}) return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( @@ -11444,6 +11377,19 @@ async def async_queue_request( ) +@router.get( + "/queue/info", + tags=["experimental"], + dependencies=[Depends(user_api_key_auth)], +) +async def queue_info( + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> List: + """Help user know the status of an item in the queue""" + return scheduler.get_queue_status() + + @router.get( "/ollama_logs", dependencies=[Depends(user_api_key_auth)], tags=["experimental"] ) diff --git a/litellm/proxy/queue/scheduler.py b/litellm/proxy/queue/scheduler.py new file mode 100644 index 00000000000..79a282d96ac --- /dev/null +++ b/litellm/proxy/queue/scheduler.py @@ -0,0 +1,129 @@ +import heapq, time +from pydantic import BaseModel +from typing import Optional +import enum +from litellm.caching import DualCache +from litellm import Router +from litellm import print_verbose + + +class SchedulerCacheKeys(enum.Enum): + queue = "scheduler:queue" + + +class DefaultPriorities(enum.Enum): + High = 0 + Medium = 128 + Low = 255 + + +class FlowItem(BaseModel): + priority: int # Priority between 0 and 255 + request_id: str + model_group: str + + +class Scheduler: + cache: DualCache + llm_router: Optional[Router] = None + + def __init__(self): + self.queue = [] + self.cache = DualCache() + + def update_variables(self, llm_router: Router, cache: Optional[DualCache] = None): + self.llm_router = llm_router + if cache is not None: + self.cache = cache + + async def add_request(self, request: FlowItem): + # We use the priority directly, as lower values indicate higher priority + # get the queue + queue = await self.get_queue(model_group=request.model_group) + # update the queue + heapq.heappush(queue, (request.priority, request.request_id)) + + # save the queue + await self.save_queue(queue=queue, model_group=request.model_group) + + async def poll(self, id: str, model_group: str) -> bool: + """Return if the id is at the top of the queue and if the token bucket allows processing""" + queue = await self.get_queue(model_group=model_group) + if not queue or not self.llm_router: + raise Exception( + "Incorrectly setup. Queue or Router is invalid. Queue={}, Router={}".format( + queue, self.llm_router + ) + ) + + # ------------ + # Setup values + # ------------ + _healthy_deployments = await self.llm_router._async_get_healthy_deployments( + model=model_group + ) + + print_verbose(f"len(_healthy_deployments): {len(_healthy_deployments)}") + if len(_healthy_deployments) == 0: + return False + + print_verbose(f"queue: {queue}, seeking id={id}") + # Check if the id is at the top of the heap + if queue[0][1] == id: + # Remove the item from the queue + heapq.heappop(queue) + print_verbose(f"Popped id: {id}") + return True + + return False + + async def peek(self, id: str, model_group: str) -> bool: + """Return if the id is at the top of the queue. Don't pop the value from heap.""" + queue = await self.get_queue(model_group=model_group) + if not queue or not self.llm_router: + raise Exception( + "Incorrectly setup. Queue or Router is invalid. Queue={}, Router={}".format( + queue, self.llm_router + ) + ) + + # ------------ + # Setup values + # ------------ + _healthy_deployments = await self.llm_router._async_get_healthy_deployments( + model=model_group + ) + if len(_healthy_deployments) == 0: + return False + + # Check if the id is at the top of the heap + if queue[0][1] == id: + return True + + return False + + def get_queue_status(self): + """Get the status of items in the queue""" + return self.queue + + async def get_queue(self, model_group: str) -> list: + """ + Return a queue for that specific model group + """ + if self.cache is not None: + _cache_key = "{}:{}".format(SchedulerCacheKeys.queue.value, model_group) + response = await self.cache.async_get_cache(key=_cache_key) + if response is None or not isinstance(response, list): + return [] + elif isinstance(response, list): + return response + return self.queue + + async def save_queue(self, queue: list, model_group: str) -> None: + """ + Save the updated queue of the model group + """ + if self.cache is not None: + _cache_key = "{}:{}".format(SchedulerCacheKeys.queue.value, model_group) + await self.cache.async_set_cache(key=_cache_key, value=queue) + return None diff --git a/litellm/tests/test_scheduler.py b/litellm/tests/test_scheduler.py new file mode 100644 index 00000000000..1177bc12c96 --- /dev/null +++ b/litellm/tests/test_scheduler.py @@ -0,0 +1,164 @@ +# What is this? +## Unit tests for the Scheduler.py (workload prioritization scheduler) + +import sys, os, time, openai, uuid +import traceback, asyncio +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from litellm import Router +from litellm.proxy.queue.scheduler import FlowItem, Scheduler + + +@pytest.mark.asyncio +async def test_scheduler_diff_model_groups(): + """ + Assert 2 requests to 2 diff model groups are top of their respective queue's + """ + scheduler = Scheduler() + + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + }, + }, + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}, + ] + ) + + scheduler.update_variables(llm_router=router) + + item1 = FlowItem(priority=0, request_id="10", model_group="gpt-3.5-turbo") + item2 = FlowItem(priority=0, request_id="11", model_group="gpt-4") + await scheduler.add_request(item1) + await scheduler.add_request(item2) + + assert await scheduler.poll(id="10", model_group="gpt-3.5-turbo") == True + assert await scheduler.poll(id="11", model_group="gpt-4") == True + + +@pytest.mark.parametrize("p0, p1", [(0, 0), (0, 1), (1, 0)]) +@pytest.mark.asyncio +async def test_scheduler_prioritized_requests(p0, p1): + """ + 2 requests for same model group + """ + scheduler = Scheduler() + + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + }, + }, + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}, + ] + ) + + scheduler.update_variables(llm_router=router) + + item1 = FlowItem(priority=p0, request_id="10", model_group="gpt-3.5-turbo") + item2 = FlowItem(priority=p1, request_id="11", model_group="gpt-3.5-turbo") + await scheduler.add_request(item1) + await scheduler.add_request(item2) + + if p0 == 0: + assert await scheduler.peek(id="10", model_group="gpt-3.5-turbo") == True + assert await scheduler.peek(id="11", model_group="gpt-3.5-turbo") == False + else: + assert await scheduler.peek(id="11", model_group="gpt-3.5-turbo") == True + assert await scheduler.peek(id="10", model_group="gpt-3.5-turbo") == False + + +@pytest.mark.parametrize("p0, p1", [(0, 0), (0, 1), (1, 0)]) +@pytest.mark.asyncio +async def test_scheduler_prioritized_requests_mock_response(p0, p1): + """ + 2 requests for same model group + """ + scheduler = Scheduler() + + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "mock_response": "Hello world this is Macintosh!", + }, + }, + ], + timeout=2, + ) + + scheduler.update_variables(llm_router=router) + + async def _make_prioritized_call(flow_item: FlowItem): + ## POLL QUEUE + default_timeout = router.timeout + end_time = time.time() + default_timeout + poll_interval = 0.03 # poll every 3ms + curr_time = time.time() + + make_request = False + + if router is None: + raise Exception("No llm router value") + + while curr_time < end_time: + make_request = await scheduler.poll( + id=flow_item.request_id, model_group=flow_item.model_group + ) + if make_request: ## IF TRUE -> MAKE REQUEST + break + else: ## ELSE -> loop till default_timeout + await asyncio.sleep(poll_interval) + curr_time = time.time() + + if make_request: + _response = await router.acompletion( + model=flow_item.model_group, + messages=[{"role": "user", "content": "Hey!"}], + ) + + return flow_item.priority, flow_item.request_id, time.time() + + raise Exception("didn't make request") + + tasks = [] + + item = FlowItem( + priority=p0, request_id=str(uuid.uuid4()), model_group="gpt-3.5-turbo" + ) + await scheduler.add_request(request=item) + tasks.append(_make_prioritized_call(flow_item=item)) + + item = FlowItem( + priority=p1, request_id=str(uuid.uuid4()), model_group="gpt-3.5-turbo" + ) + await scheduler.add_request(request=item) + tasks.append(_make_prioritized_call(flow_item=item)) + + # Running the tasks and getting responses in order of completion + completed_responses = [] + for task in asyncio.as_completed(tasks): + result = await task + completed_responses.append(result) + print(f"Received response: {result}") + + print(f"responses: {completed_responses}") + assert ( + completed_responses[0][0] == 0 + ) # assert higher priority request got done first + assert ( + completed_responses[0][2] < completed_responses[1][2] + ), "1st response time={}, 2nd response time={}".format( + completed_responses[0][1], completed_responses[1][1] + ) # assert higher priority request got done first diff --git a/litellm/types/router.py b/litellm/types/router.py index 75e792f4c92..f131c29298b 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -326,6 +326,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): output_cost_per_token: Optional[float] input_cost_per_second: Optional[float] output_cost_per_second: Optional[float] + ## MOCK RESPONSES ## + mock_response: Optional[str] class DeploymentTypedDict(TypedDict): diff --git a/tests/test_openai_endpoints.py b/tests/test_openai_endpoints.py index 43dcae3cd7b..83d387ffbec 100644 --- a/tests/test_openai_endpoints.py +++ b/tests/test_openai_endpoints.py @@ -103,6 +103,36 @@ async def chat_completion(session, key, model: Union[str, List] = "gpt-4"): return await response.json() +async def queue_chat_completion( + session, key, priority: int, model: Union[str, List] = "gpt-4" +): + url = "http://0.0.0.0:4000/queue/chat/completions" + headers = { + "Authorization": f"Bearer {key}", + "Content-Type": "application/json", + } + data = { + "model": model, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"}, + ], + "priority": priority, + } + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(response_text) + print() + + if status != 200: + raise Exception(f"Request did not return a 200 status code: {status}") + + return response.raw_headers + + async def chat_completion_with_headers(session, key, model="gpt-4"): url = "http://0.0.0.0:4000/chat/completions" headers = { From f8d4be710ec2b235ef421d0491ac390da9be34b4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 31 May 2024 19:35:47 -0700 Subject: [PATCH 2/8] docs(scheduler.md): add request prioritization to docs --- docs/my-website/docs/scheduler.md | 141 +++++++++++++++++++++++++ docs/my-website/sidebars.js | 1 + litellm/__init__.py | 1 + litellm/proxy/proxy_server.py | 8 +- litellm/{proxy/queue => }/scheduler.py | 58 +++++----- litellm/tests/test_scheduler.py | 49 ++++----- 6 files changed, 205 insertions(+), 53 deletions(-) create mode 100644 docs/my-website/docs/scheduler.md rename litellm/{proxy/queue => }/scheduler.py (70%) diff --git a/docs/my-website/docs/scheduler.md b/docs/my-website/docs/scheduler.md new file mode 100644 index 00000000000..347406ade25 --- /dev/null +++ b/docs/my-website/docs/scheduler.md @@ -0,0 +1,141 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# [BETA] Request Prioritization + +:::info + +Beta feature. Use for testing only. + +[Help us improve this](https://github.com/BerriAI/litellm/issues) +::: + +Prioritize LLM API requests in high-traffic. + +- Add request to priority queue +- Poll queue, to check if request can be made. Returns 'True': + * if there's healthy deployments + * OR if request is at top of queue +- Priority - The lower the number, the higher the priority: + * e.g. `priority=0` > `priority=2000` + +## Quick Start + +```python +from litellm import Scheduler, FlowItem, Router + +scheduler = Scheduler() + +router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "mock_response": "Hello world this is Macintosh!", # fakes the LLM API call + "rpm": 1, + }, + }, + ], + timeout=2, # timeout request if takes > 2s + routing_strategy="usage-based-routing-v2", +) + +scheduler.update_variables(llm_router=router) + +### 🚨 IMPORTANT ### + +item = FlowItem( + priority=0, # 👈 SET PRIORITY FOR REQUEST + request_id=str(uuid.uuid4()), # 👈 SET REQUEST ID + model_name="gpt-3.5-turbo" # 👈 SAME as 'Router' +) + +### [fin] IMPORTANT ### + +## ADDS REQUEST TO QUEUE ## +await scheduler.add_request(request=item) + +## POLL QUEUE +default_timeout = router.timeout +end_time = time.time() + default_timeout +poll_interval = 0.03 # poll every 3ms +curr_time = time.time() + +make_request = False + +while curr_time < end_time: + make_request = await 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 + ) + if make_request: ## IF TRUE -> MAKE REQUEST + break + else: ## ELSE -> loop till default_timeout + await asyncio.sleep(poll_interval) + curr_time = time.time() + +if make_request: + try: + _response = await router.acompletion( + model=item.model_name, + messages=[{"role": "user", "content": "Hey!"}], + ) + except Exception as e: + print("{}, {}, {}".format(item.priority, item.request_id, "Error occurred")) + + print("{}, {}, {}".format(item.priority, item.request_id, time.time())) + +print("didn't make request") +``` + +## LiteLLM Proxy + +To prioritize requests on LiteLLM Proxy call our beta openai-compatible `http://localhost:4000/queue` endpoint. + + + + +```curl +curl -X POST 'http://localhost:4000/queue/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-D '{ + "model": "gpt-3.5-turbo-fake-model", + "messages": [ + { + "role": "user", + "content": "what is the meaning of the universe? 1234" + }], + "priority": 0 👈 SET VALUE HERE +}' +``` + + + + +```python +import openai +client = openai.OpenAI( + api_key="anything", + base_url="http://0.0.0.0:4000" +) + +# request sent to model set on litellm proxy, `litellm --model` +response = client.chat.completions.create( + model="gpt-3.5-turbo", + messages = [ + { + "role": "user", + "content": "this is a test request, write a short poem" + } + ], + extra_body={ + "priority": 0 👈 SET VALUE HERE + } +) + +print(response) +``` + + + \ No newline at end of file diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 29095d41fe9..d5c6f8bc288 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -164,6 +164,7 @@ const sidebars = { }, "proxy/custom_pricing", "routing", + "scheduler", "rules", "set_keys", "budget_manager", diff --git a/litellm/__init__.py b/litellm/__init__.py index 9fa801318e1..6097b8c94a3 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -805,3 +805,4 @@ from .proxy.proxy_cli import run_server from .router import Router from .assistants.main import * from .batches.main import * +from .scheduler import * diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3a0efd653f3..f5289b02484 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -141,7 +141,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.exceptions import RejectedRequestError from litellm.integrations.slack_alerting import SlackAlertingArgs, SlackAlerting -from litellm.proxy.queue.scheduler import Scheduler, FlowItem, DefaultPriorities +from litellm.scheduler import Scheduler, FlowItem, DefaultPriorities try: from litellm._version import version @@ -11305,7 +11305,7 @@ async def async_queue_request( flow_item = FlowItem( priority=data.pop("priority", DefaultPriorities.Medium.value), request_id=request_id, - model_group=data["model"], + model_name=data["model"], ) # [TODO] only allow premium users to set non default priorities @@ -11330,9 +11330,7 @@ async def async_queue_request( ) while curr_time < end_time: - make_request = await scheduler.poll( - id=request_id, model_group=data["model"] - ) + make_request = await scheduler.poll(id=request_id, model_name=data["model"]) if make_request: ## IF TRUE -> MAKE REQUEST break else: ## ELSE -> loop till default_timeout diff --git a/litellm/proxy/queue/scheduler.py b/litellm/scheduler.py similarity index 70% rename from litellm/proxy/queue/scheduler.py rename to litellm/scheduler.py index 79a282d96ac..3bbd3916e60 100644 --- a/litellm/proxy/queue/scheduler.py +++ b/litellm/scheduler.py @@ -20,7 +20,7 @@ class DefaultPriorities(enum.Enum): class FlowItem(BaseModel): priority: int # Priority between 0 and 255 request_id: str - model_group: str + model_name: str class Scheduler: @@ -39,16 +39,26 @@ class Scheduler: async def add_request(self, request: FlowItem): # We use the priority directly, as lower values indicate higher priority # get the queue - queue = await self.get_queue(model_group=request.model_group) + queue = await self.get_queue(model_name=request.model_name) # update the queue heapq.heappush(queue, (request.priority, request.request_id)) # save the queue - await self.save_queue(queue=queue, model_group=request.model_group) + await self.save_queue(queue=queue, model_name=request.model_name) - async def poll(self, id: str, model_group: str) -> bool: - """Return if the id is at the top of the queue and if the token bucket allows processing""" - queue = await self.get_queue(model_group=model_group) + async def poll(self, id: str, model_name: str) -> bool: + """ + Return if request can be processed. + + Returns: + - True: + * If healthy deployments are available + * OR If request at the top of queue + - False: + * If no healthy deployments available + * AND request not at the top of queue + """ + queue = await self.get_queue(model_name=model_name) if not queue or not self.llm_router: raise Exception( "Incorrectly setup. Queue or Router is invalid. Queue={}, Router={}".format( @@ -60,26 +70,26 @@ class Scheduler: # Setup values # ------------ _healthy_deployments = await self.llm_router._async_get_healthy_deployments( - model=model_group + model=model_name ) print_verbose(f"len(_healthy_deployments): {len(_healthy_deployments)}") if len(_healthy_deployments) == 0: - return False + print_verbose(f"queue: {queue}, seeking id={id}") + # Check if the id is at the top of the heap + if queue[0][1] == id: + # Remove the item from the queue + heapq.heappop(queue) + print_verbose(f"Popped id: {id}") + return True + else: + return False - print_verbose(f"queue: {queue}, seeking id={id}") - # Check if the id is at the top of the heap - if queue[0][1] == id: - # Remove the item from the queue - heapq.heappop(queue) - print_verbose(f"Popped id: {id}") - return True + return True - return False - - async def peek(self, id: str, model_group: str) -> bool: + async def peek(self, id: str, model_name: str) -> bool: """Return if the id is at the top of the queue. Don't pop the value from heap.""" - queue = await self.get_queue(model_group=model_group) + queue = await self.get_queue(model_name=model_name) if not queue or not self.llm_router: raise Exception( "Incorrectly setup. Queue or Router is invalid. Queue={}, Router={}".format( @@ -91,7 +101,7 @@ class Scheduler: # Setup values # ------------ _healthy_deployments = await self.llm_router._async_get_healthy_deployments( - model=model_group + model=model_name ) if len(_healthy_deployments) == 0: return False @@ -106,12 +116,12 @@ class Scheduler: """Get the status of items in the queue""" return self.queue - async def get_queue(self, model_group: str) -> list: + async def get_queue(self, model_name: str) -> list: """ Return a queue for that specific model group """ if self.cache is not None: - _cache_key = "{}:{}".format(SchedulerCacheKeys.queue.value, model_group) + _cache_key = "{}:{}".format(SchedulerCacheKeys.queue.value, model_name) response = await self.cache.async_get_cache(key=_cache_key) if response is None or not isinstance(response, list): return [] @@ -119,11 +129,11 @@ class Scheduler: return response return self.queue - async def save_queue(self, queue: list, model_group: str) -> None: + async def save_queue(self, queue: list, model_name: str) -> None: """ Save the updated queue of the model group """ if self.cache is not None: - _cache_key = "{}:{}".format(SchedulerCacheKeys.queue.value, model_group) + _cache_key = "{}:{}".format(SchedulerCacheKeys.queue.value, model_name) await self.cache.async_set_cache(key=_cache_key, value=queue) return None diff --git a/litellm/tests/test_scheduler.py b/litellm/tests/test_scheduler.py index 1177bc12c96..2e48eab3c46 100644 --- a/litellm/tests/test_scheduler.py +++ b/litellm/tests/test_scheduler.py @@ -9,11 +9,11 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path from litellm import Router -from litellm.proxy.queue.scheduler import FlowItem, Scheduler +from litellm.scheduler import FlowItem, Scheduler @pytest.mark.asyncio -async def test_scheduler_diff_model_groups(): +async def test_scheduler_diff_model_names(): """ Assert 2 requests to 2 diff model groups are top of their respective queue's """ @@ -33,13 +33,13 @@ async def test_scheduler_diff_model_groups(): scheduler.update_variables(llm_router=router) - item1 = FlowItem(priority=0, request_id="10", model_group="gpt-3.5-turbo") - item2 = FlowItem(priority=0, request_id="11", model_group="gpt-4") + item1 = FlowItem(priority=0, request_id="10", model_name="gpt-3.5-turbo") + item2 = FlowItem(priority=0, request_id="11", model_name="gpt-4") await scheduler.add_request(item1) await scheduler.add_request(item2) - assert await scheduler.poll(id="10", model_group="gpt-3.5-turbo") == True - assert await scheduler.poll(id="11", model_group="gpt-4") == True + assert await scheduler.poll(id="10", model_name="gpt-3.5-turbo") == True + assert await scheduler.poll(id="11", model_name="gpt-4") == True @pytest.mark.parametrize("p0, p1", [(0, 0), (0, 1), (1, 0)]) @@ -64,17 +64,17 @@ async def test_scheduler_prioritized_requests(p0, p1): scheduler.update_variables(llm_router=router) - item1 = FlowItem(priority=p0, request_id="10", model_group="gpt-3.5-turbo") - item2 = FlowItem(priority=p1, request_id="11", model_group="gpt-3.5-turbo") + item1 = FlowItem(priority=p0, request_id="10", model_name="gpt-3.5-turbo") + item2 = FlowItem(priority=p1, request_id="11", model_name="gpt-3.5-turbo") await scheduler.add_request(item1) await scheduler.add_request(item2) if p0 == 0: - assert await scheduler.peek(id="10", model_group="gpt-3.5-turbo") == True - assert await scheduler.peek(id="11", model_group="gpt-3.5-turbo") == False + assert await scheduler.peek(id="10", model_name="gpt-3.5-turbo") == True + assert await scheduler.peek(id="11", model_name="gpt-3.5-turbo") == False else: - assert await scheduler.peek(id="11", model_group="gpt-3.5-turbo") == True - assert await scheduler.peek(id="10", model_group="gpt-3.5-turbo") == False + assert await scheduler.peek(id="11", model_name="gpt-3.5-turbo") == True + assert await scheduler.peek(id="10", model_name="gpt-3.5-turbo") == False @pytest.mark.parametrize("p0, p1", [(0, 0), (0, 1), (1, 0)]) @@ -92,10 +92,12 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): "litellm_params": { "model": "gpt-3.5-turbo", "mock_response": "Hello world this is Macintosh!", + "rpm": 1, }, }, ], timeout=2, + routing_strategy="usage-based-routing-v2", ) scheduler.update_variables(llm_router=router) @@ -114,7 +116,7 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): while curr_time < end_time: make_request = await scheduler.poll( - id=flow_item.request_id, model_group=flow_item.model_group + id=flow_item.request_id, model_name=flow_item.model_name ) if make_request: ## IF TRUE -> MAKE REQUEST break @@ -123,10 +125,13 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): curr_time = time.time() if make_request: - _response = await router.acompletion( - model=flow_item.model_group, - messages=[{"role": "user", "content": "Hey!"}], - ) + try: + _response = await router.acompletion( + model=flow_item.model_name, + messages=[{"role": "user", "content": "Hey!"}], + ) + except Exception as e: + return flow_item.priority, flow_item.request_id, "Error occurred" return flow_item.priority, flow_item.request_id, time.time() @@ -135,13 +140,13 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): tasks = [] item = FlowItem( - priority=p0, request_id=str(uuid.uuid4()), model_group="gpt-3.5-turbo" + priority=p0, request_id=str(uuid.uuid4()), model_name="gpt-3.5-turbo" ) await scheduler.add_request(request=item) tasks.append(_make_prioritized_call(flow_item=item)) item = FlowItem( - priority=p1, request_id=str(uuid.uuid4()), model_group="gpt-3.5-turbo" + priority=p1, request_id=str(uuid.uuid4()), model_name="gpt-3.5-turbo" ) await scheduler.add_request(request=item) tasks.append(_make_prioritized_call(flow_item=item)) @@ -157,8 +162,4 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): assert ( completed_responses[0][0] == 0 ) # assert higher priority request got done first - assert ( - completed_responses[0][2] < completed_responses[1][2] - ), "1st response time={}, 2nd response time={}".format( - completed_responses[0][1], completed_responses[1][1] - ) # assert higher priority request got done first + assert isinstance(completed_responses[1][2], str) # 2nd request errored out From f3c37f487aa02716f69b48f72878e7255e2d3bfe Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 31 May 2024 19:41:53 -0700 Subject: [PATCH 3/8] fix(proxy_server.py): fix naming conflict --- litellm/proxy/proxy_server.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f5289b02484..77e781a0830 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -397,7 +397,7 @@ async_result = None celery_app_conn = None celery_fn = None # Redis Queue for handling requests ### SIMPLE QUEUE ### -scheduler = Scheduler() +simple_scheduler = Scheduler() ### DB WRITER ### db_writer_client: Optional[HTTPHandler] = None ### logger ### @@ -3658,7 +3658,7 @@ def on_backoff(details): @router.on_event("startup") async def startup_event(): - global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, scheduler + global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, simple_scheduler import json ### LOAD MASTER KEY ### @@ -3696,7 +3696,7 @@ async def startup_event(): ## Priority Workload Scheduler ## if llm_router is not None: - scheduler.update_variables(llm_router=llm_router) + simple_scheduler.update_variables(llm_router=llm_router) ## UPDATE SLACK ALERTING ## proxy_logging_obj.slack_alerting_instance.update_values(llm_router=llm_router) @@ -11310,7 +11310,7 @@ async def async_queue_request( # [TODO] only allow premium users to set non default priorities ## ADD REQUEST TO QUEUE - response = await scheduler.add_request(request=flow_item) + response = await simple_scheduler.add_request(request=flow_item) if llm_router is None: raise HTTPException( @@ -11330,7 +11330,9 @@ async def async_queue_request( ) while curr_time < end_time: - make_request = await scheduler.poll(id=request_id, model_name=data["model"]) + make_request = await simple_scheduler.poll( + id=request_id, model_name=data["model"] + ) if make_request: ## IF TRUE -> MAKE REQUEST break else: ## ELSE -> loop till default_timeout @@ -11385,7 +11387,7 @@ async def queue_info( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> List: """Help user know the status of an item in the queue""" - return scheduler.get_queue_status() + return simple_scheduler.get_queue_status() @router.get( From e49325b2348a9013ad4553131cd39cfb8aa36565 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 31 May 2024 21:32:01 -0700 Subject: [PATCH 4/8] fix(router.py): fix cooldown logic for usage-based-routing-v2 pre-call-checks --- litellm/main.py | 5 +++++ litellm/router.py | 28 ++++++++++++++++++++++------ litellm/tests/test_scheduler.py | 22 ++++++++++++++++++---- 3 files changed, 45 insertions(+), 10 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 525a39d6899..32cad89af80 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -430,6 +430,10 @@ def mock_completion( model=model, # type: ignore request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), ) + time_delay = kwargs.get("mock_delay", None) + if time_delay is not None: + time.sleep(time_delay) + model_response = ModelResponse(stream=stream) if stream is True: # don't try to access stream object, @@ -880,6 +884,7 @@ def completion( mock_response=mock_response, logging=logging, acompletion=acompletion, + mock_delay=kwargs.get("mock_delay", None), ) if custom_llm_provider == "azure": # azure configs diff --git a/litellm/router.py b/litellm/router.py index d7ff84877c5..c0abad448f8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -631,7 +631,6 @@ class Router: kwargs=kwargs, client_type="max_parallel_requests", ) - if rpm_semaphore is not None and isinstance( rpm_semaphore, asyncio.Semaphore ): @@ -1875,6 +1874,7 @@ class Router: error=e, healthy_deployments=_healthy_deployments, context_window_fallbacks=context_window_fallbacks, + regular_fallbacks=fallbacks, ) # decides how long to sleep before retry @@ -1884,7 +1884,6 @@ class Router: num_retries=num_retries, healthy_deployments=_healthy_deployments, ) - # sleeps for the length of the timeout await asyncio.sleep(_timeout) @@ -1929,6 +1928,7 @@ class Router: healthy_deployments=_healthy_deployments, ) await asyncio.sleep(_timeout) + try: cooldown_deployments = await self._async_get_cooldown_deployments() original_exception.message += f"\nNumber Retries = {current_attempt + 1}, Max Retries={num_retries}\nCooldown Deployments={cooldown_deployments}" @@ -1941,6 +1941,7 @@ class Router: error: Exception, healthy_deployments: Optional[List] = None, context_window_fallbacks: Optional[List] = None, + regular_fallbacks: Optional[List] = None, ): """ 1. raise an exception for ContextWindowExceededError if context_window_fallbacks is not None @@ -1957,7 +1958,7 @@ class Router: ### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR w/ fallbacks available / Bad Request Error if ( isinstance(error, litellm.ContextWindowExceededError) - and context_window_fallbacks is None + and context_window_fallbacks is not None ): raise error @@ -1965,7 +1966,11 @@ class Router: if isinstance(error, openai.RateLimitError) or isinstance( error, openai.AuthenticationError ): - if _num_healthy_deployments <= 0: + if ( + _num_healthy_deployments <= 0 + and regular_fallbacks is not None + and len(regular_fallbacks) > 0 + ): raise error return True @@ -2140,6 +2145,7 @@ class Router: error=e, healthy_deployments=_healthy_deployments, context_window_fallbacks=context_window_fallbacks, + fallbacks=fallbacks, ) # decides how long to sleep before retry @@ -2348,7 +2354,7 @@ class Router: the exception is not one that should be immediately retried (e.g. 401) """ - args = locals() + if deployment is None: return @@ -2519,7 +2525,17 @@ class Router: """ for _callback in litellm.callbacks: if isinstance(_callback, CustomLogger): - response = await _callback.async_pre_call_check(deployment) + try: + response = await _callback.async_pre_call_check(deployment) + except litellm.RateLimitError as e: + self._set_cooldown_deployments( + exception_status=e.status_code, + deployment=deployment["model_info"]["id"], + time_to_cooldown=self.cooldown_time, + ) + raise e + except Exception as e: + raise e def set_client(self, model: dict): """ diff --git a/litellm/tests/test_scheduler.py b/litellm/tests/test_scheduler.py index 2e48eab3c46..bba06d5873d 100644 --- a/litellm/tests/test_scheduler.py +++ b/litellm/tests/test_scheduler.py @@ -77,11 +77,13 @@ async def test_scheduler_prioritized_requests(p0, p1): assert await scheduler.peek(id="10", model_name="gpt-3.5-turbo") == False -@pytest.mark.parametrize("p0, p1", [(0, 0), (0, 1), (1, 0)]) +@pytest.mark.parametrize("p0, p1", [(0, 1)]) # (0, 0), (1, 0) @pytest.mark.asyncio async def test_scheduler_prioritized_requests_mock_response(p0, p1): """ 2 requests for same model group + + if model is at rate limit, ensure the higher priority request gets done first """ scheduler = Scheduler() @@ -96,12 +98,19 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): }, }, ], - timeout=2, + timeout=10, + num_retries=3, + cooldown_time=5, routing_strategy="usage-based-routing-v2", ) scheduler.update_variables(llm_router=router) + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey!"}], + ) + async def _make_prioritized_call(flow_item: FlowItem): ## POLL QUEUE default_timeout = router.timeout @@ -118,6 +127,7 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): make_request = await scheduler.poll( id=flow_item.request_id, model_name=flow_item.model_name ) + print(f"make_request={make_request}, priority={flow_item.priority}") if make_request: ## IF TRUE -> MAKE REQUEST break else: ## ELSE -> loop till default_timeout @@ -131,7 +141,8 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): messages=[{"role": "user", "content": "Hey!"}], ) except Exception as e: - return flow_item.priority, flow_item.request_id, "Error occurred" + print("Received error - {}".format(str(e))) + return flow_item.priority, flow_item.request_id, time.time() return flow_item.priority, flow_item.request_id, time.time() @@ -159,7 +170,10 @@ async def test_scheduler_prioritized_requests_mock_response(p0, p1): print(f"Received response: {result}") print(f"responses: {completed_responses}") + assert ( completed_responses[0][0] == 0 ) # assert higher priority request got done first - assert isinstance(completed_responses[1][2], str) # 2nd request errored out + assert ( + completed_responses[0][2] < completed_responses[1][2] + ) # higher priority request tried first From 8df912cfa2ec7a93e67efcbcfeb2ca2448218195 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 31 May 2024 21:42:00 -0700 Subject: [PATCH 5/8] Delete litellm/proxy/_experimental/out/model_hub/index.html --- litellm/proxy/_experimental/out/model_hub/index.html | 1 - 1 file changed, 1 deletion(-) delete mode 100644 litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub/index.html b/litellm/proxy/_experimental/out/model_hub/index.html deleted file mode 100644 index cdaa2157dc5..00000000000 --- a/litellm/proxy/_experimental/out/model_hub/index.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file From 989665c7dce814e00b9cc3c102236fab8bce7c07 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 31 May 2024 21:42:15 -0700 Subject: [PATCH 6/8] Delete litellm/proxy/_experimental/out/404/index.html --- litellm/proxy/_experimental/out/404/index.html | 1 - 1 file changed, 1 deletion(-) delete mode 100644 litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/404/index.html b/litellm/proxy/_experimental/out/404/index.html deleted file mode 100644 index a5784108815..00000000000 --- a/litellm/proxy/_experimental/out/404/index.html +++ /dev/null @@ -1 +0,0 @@ -404: This page could not be found.LiteLLM Dashboard

404

This page could not be found.

\ No newline at end of file From 381247a0959b9fcc58a32bfb390b6da14432bbfa Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 31 May 2024 21:49:30 -0700 Subject: [PATCH 7/8] fix(router.py): fix param --- litellm/router.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index c0abad448f8..10cf609a0c2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2145,7 +2145,7 @@ class Router: error=e, healthy_deployments=_healthy_deployments, context_window_fallbacks=context_window_fallbacks, - fallbacks=fallbacks, + regular_fallbacks=fallbacks, ) # decides how long to sleep before retry From 183dd81087ba522dd234c2d3a05de477abf028b4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 31 May 2024 23:27:43 -0700 Subject: [PATCH 8/8] fix(router.py): fix should_retry logic --- litellm/tests/test_router_retries.py | 43 +++++++++++++++------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py index 7273fd6e968..1c0c3e49f83 100644 --- a/litellm/tests/test_router_retries.py +++ b/litellm/tests/test_router_retries.py @@ -134,7 +134,7 @@ async def test_router_retry_policy(error_type): ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0 ) - router = litellm.Router( + router = Router( model_list=[ { "model_name": "gpt-3.5-turbo", # openai model name @@ -334,13 +334,13 @@ def test_retry_rate_limit_error_with_healthy_deployments(): ) -def test_do_not_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployments(): +def test_do_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployments(): """ - Test 2. It SHOULD NOT Retry, when healthy_deployments is [] and fallbacks is None + Test 2. It SHOULD Retry, when healthy_deployments is [] and fallbacks is None """ healthy_deployments = [] - router = litellm.Router( + router = Router( model_list=[ { "model_name": "gpt-3.5-turbo", @@ -359,14 +359,14 @@ def test_do_not_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployme response = router.should_retry_this_error( error=rate_limit_error, healthy_deployments=healthy_deployments ) - assert response != True, "Should have raised RateLimitError" - except openai.RateLimitError: - pass + assert response == True + except Exception as e: + pytest.fail("Should not have failed this error - {}".format(str(e))) def test_raise_context_window_exceeded_error(): """ - Retry Context Window Exceeded Error, when context_window_fallbacks is not None + Trigger Context Window fallback, when context_window_fallbacks is not None """ context_window_error = litellm.ContextWindowExceededError( message="Context window exceeded", @@ -379,7 +379,7 @@ def test_raise_context_window_exceeded_error(): ) context_window_fallbacks = [{"gpt-3.5-turbo": ["azure/chatgpt-v-2"]}] - router = litellm.Router( + router = Router( model_list=[ { "model_name": "gpt-3.5-turbo", @@ -393,14 +393,17 @@ def test_raise_context_window_exceeded_error(): ] ) - response = router.should_retry_this_error( - error=context_window_error, - healthy_deployments=None, - context_window_fallbacks=context_window_fallbacks, - ) - assert ( - response == True - ), "Should not have raised exception since we have context window fallbacks" + try: + response = router.should_retry_this_error( + error=context_window_error, + healthy_deployments=None, + context_window_fallbacks=context_window_fallbacks, + ) + pytest.fail( + "Expected to raise context window exceeded error -> trigger fallback" + ) + except Exception as e: + pass def test_raise_context_window_exceeded_error_no_retry(): @@ -418,7 +421,7 @@ def test_raise_context_window_exceeded_error_no_retry(): ) context_window_fallbacks = None - router = litellm.Router( + router = Router( model_list=[ { "model_name": "gpt-3.5-turbo", @@ -439,8 +442,8 @@ def test_raise_context_window_exceeded_error_no_retry(): context_window_fallbacks=context_window_fallbacks, ) assert ( - response != True - ), "Should have raised exception since we do not have context window fallbacks" + response == True + ), "Should not have raised exception since we do not have context window fallbacks" except litellm.ContextWindowExceededError: pass