From 0b729046087a3646ecf7c5573fb6d4cf861e520a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 3 May 2024 09:00:32 -0700 Subject: [PATCH 1/2] fix(lowest_latency.py): fix the size of the latency list to 10 by default (can be modified) --- litellm/router_strategy/lowest_latency.py | 37 +++++++- litellm/tests/test_lowest_latency_routing.py | 92 +++++++++++++++++++- 2 files changed, 124 insertions(+), 5 deletions(-) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 80dee5e678..5f0f15aac0 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -31,6 +31,7 @@ class LiteLLMBase(BaseModel): class RoutingArgs(LiteLLMBase): ttl: int = 1 * 60 * 60 # 1 hour lowest_latency_buffer: float = 0 + max_latency_list_size: int = 10 class LowestLatencyLoggingHandler(CustomLogger): @@ -103,7 +104,18 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - request_count_dict[id].setdefault("latency", []).append(final_value) + if ( + len(request_count_dict[id].get("latency", [])) + < self.routing_args.max_latency_list_size + ): + request_count_dict[id].setdefault("latency", []).append(final_value) + else: + request_count_dict[id]["latency"] = request_count_dict[id][ + "latency" + ][: self.routing_args.max_latency_list_size - 1] + [final_value] + + if precise_minute not in request_count_dict[id]: + request_count_dict[id][precise_minute] = {} if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -170,8 +182,17 @@ class LowestLatencyLoggingHandler(CustomLogger): if id not in request_count_dict: request_count_dict[id] = {} - ## Latency - request_count_dict[id].setdefault("latency", []).append(1000.0) + ## Latency - give 1000s penalty for failing + if ( + len(request_count_dict[id].get("latency", [])) + < self.routing_args.max_latency_list_size + ): + request_count_dict[id].setdefault("latency", []).append(1000.0) + else: + request_count_dict[id]["latency"] = request_count_dict[id][ + "latency" + ][: self.routing_args.max_latency_list_size - 1] + [1000.0] + self.router_cache.set_cache( key=latency_key, value=request_count_dict, @@ -242,7 +263,15 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - request_count_dict[id].setdefault("latency", []).append(final_value) + if ( + len(request_count_dict[id].get("latency", [])) + < self.routing_args.max_latency_list_size + ): + request_count_dict[id].setdefault("latency", []).append(final_value) + else: + request_count_dict[id]["latency"] = request_count_dict[id][ + "latency" + ][: self.routing_args.max_latency_list_size - 1] + [final_value] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} diff --git a/litellm/tests/test_lowest_latency_routing.py b/litellm/tests/test_lowest_latency_routing.py index 2f0aaee91d..4da8792087 100644 --- a/litellm/tests/test_lowest_latency_routing.py +++ b/litellm/tests/test_lowest_latency_routing.py @@ -7,7 +7,7 @@ import traceback from dotenv import load_dotenv load_dotenv() -import os +import os, copy sys.path.insert( 0, os.path.abspath("../..") @@ -20,6 +20,96 @@ from litellm.caching import DualCache ### UNIT TESTS FOR LATENCY ROUTING ### +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_latency_memory_leak(sync_mode): + """ + Test to make sure there's no memory leak caused by lowest latency routing + + - make 10 calls -> check memory + - make 11th call -> no change in memory + """ + test_cache = DualCache() + model_list = [] + lowest_latency_logger = LowestLatencyLoggingHandler( + router_cache=test_cache, model_list=model_list + ) + model_group = "gpt-3.5-turbo" + deployment_id = "1234" + kwargs = { + "litellm_params": { + "metadata": { + "model_group": "gpt-3.5-turbo", + "deployment": "azure/chatgpt-v-2", + }, + "model_info": {"id": deployment_id}, + } + } + start_time = time.time() + response_obj = {"usage": {"total_tokens": 50}} + time.sleep(5) + end_time = time.time() + for _ in range(10): + if sync_mode: + lowest_latency_logger.log_success_event( + response_obj=response_obj, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + else: + await lowest_latency_logger.async_log_success_event( + response_obj=response_obj, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + latency_key = f"{model_group}_map" + cache_value = copy.deepcopy( + test_cache.get_cache(key=latency_key) + ) # MAKE SURE NO MEMORY LEAK IN CACHING OBJECT + + if sync_mode: + lowest_latency_logger.log_success_event( + response_obj=response_obj, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + else: + await lowest_latency_logger.async_log_success_event( + response_obj=response_obj, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + new_cache_value = test_cache.get_cache(key=latency_key) + # Assert that the size of the cache doesn't grow unreasonably + assert get_size(new_cache_value) <= get_size( + cache_value + ), f"Memory leak detected in function call! new_cache size={get_size(new_cache_value)}, old cache size={get_size(cache_value)}" + + +def get_size(obj, seen=None): + # From https://goshippo.com/blog/measure-real-size-any-python-object/ + # Recursively finds size of objects + size = sys.getsizeof(obj) + if seen is None: + seen = set() + obj_id = id(obj) + if obj_id in seen: + return 0 + seen.add(obj_id) + if isinstance(obj, dict): + size += sum([get_size(v, seen) for v in obj.values()]) + size += sum([get_size(k, seen) for k in obj.keys()]) + elif hasattr(obj, "__dict__"): + size += get_size(obj.__dict__, seen) + elif hasattr(obj, "__iter__") and not isinstance(obj, (str, bytes, bytearray)): + size += sum([get_size(i, seen) for i in obj]) + return size + + def test_latency_updated(): test_cache = DualCache() model_list = [] From 2dd9d2f704028be562f7fd1cbd4709300e3f5c47 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 3 May 2024 10:09:57 -0700 Subject: [PATCH 2/2] test(test_amazing_vertex_completion.py): try-except api errors --- .../tests/test_amazing_vertex_completion.py | 36 ------------------- litellm/utils.py | 21 +++++++---- 2 files changed, 15 insertions(+), 42 deletions(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 05eece8344..1d79653ea6 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -548,42 +548,6 @@ def test_gemini_pro_vision_base64(): def test_gemini_pro_function_calling(): - load_vertex_ai_credentials() - tools = [ - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - }, - "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, - }, - "required": ["location"], - }, - }, - } - ] - - messages = [ - { - "role": "user", - "content": "What's the weather like in Boston today in fahrenheit?", - } - ] - completion = litellm.completion( - model="gemini-pro", messages=messages, tools=tools, tool_choice="auto" - ) - print(f"completion: {completion}") - if hasattr(completion.choices[0].message, "tool_calls") and isinstance( - completion.choices[0].message.tool_calls, list - ): - assert len(completion.choices[0].message.tool_calls) == 1 try: load_vertex_ai_credentials() tools = [ diff --git a/litellm/utils.py b/litellm/utils.py index ec296e9dc3..80d26f58b9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3974,12 +3974,10 @@ def calculage_img_tokens( tile_tokens = (base_tokens * 2) * tiles_needed_high_res total_tokens = base_tokens + tile_tokens return total_tokens - + def create_pretrained_tokenizer( - identifier: str, - revision="main", - auth_token: Optional[str] = None + identifier: str, revision="main", auth_token: Optional[str] = None ): """ Creates a tokenizer from an existing file on a HuggingFace repository to be used with `token_counter`. @@ -3993,7 +3991,9 @@ def create_pretrained_tokenizer( dict: A dictionary with the tokenizer and its type. """ - tokenizer = Tokenizer.from_pretrained(identifier, revision=revision, auth_token=auth_token) + tokenizer = Tokenizer.from_pretrained( + identifier, revision=revision, auth_token=auth_token + ) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} @@ -9001,7 +9001,16 @@ def exception_type( request=original_exception.request, ) elif custom_llm_provider == "azure": - if "This model's maximum context length is" in error_str: + if "Internal server error" in error_str: + exception_mapping_worked = True + raise APIError( + status_code=500, + message=f"AzureException - {original_exception.message}", + llm_provider="azure", + model=model, + request=httpx.Request(method="POST", url="https://openai.com/"), + ) + elif "This model's maximum context length is" in error_str: exception_mapping_worked = True raise ContextWindowExceededError( message=f"AzureException - {original_exception.message}",