Merge pull request #3422 from BerriAI/litellm_lowest_latency_fix

fix(lowest_latency.py): fix the size of the latency list to 10 by default (can be modified)

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] = {}

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 = [

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 = []

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}",