Litellm fix multi instance checks on teams (#11137)

* fix(auth_checks.py): set ttl on in-memory team objects

ensures objects expire

* fix(in_memory_cache.py): allow ttl override if existing ttl expired (no 'get_cache' ever called)

litellm/caching/in_memory_cache.py

@@ -106,6 +106,18 @@ class InMemoryCache(BaseCache):
                 # One of the most common causes of memory leaks in Python is the retention of objects that are no longer being used.
                 # This can occur when an object is referenced by another object, but the reference is never removed.
 
+    def allow_ttl_override(self, key: str) -> bool:
+        """
+        Check if ttl is set for a key
+        """
+        ttl_time = self.ttl_dict.get(key)
+        if ttl_time is None:  # if ttl is not set, allow override
+            return True
+        elif float(ttl_time) < time.time():  # if ttl is expired, allow override
+            return True
+        else:
+            return False
+
     def set_cache(self, key, value, **kwargs):
         if len(self.cache_dict) >= self.max_size_in_memory:
             # only evict when cache is full
@@ -114,7 +126,7 @@ class InMemoryCache(BaseCache):
             return
 
         self.cache_dict[key] = value
-        if self.ttl_dict.get(key) is None:  # if ttl is not set, set it to default ttl
+        if self.allow_ttl_override(key):  # if ttl is not set, set it to default ttl
             if "ttl" in kwargs and kwargs["ttl"] is not None:
                 self.ttl_dict[key] = time.time() + kwargs["ttl"]
             else:

litellm/proxy/_new_secret_config.yaml

@@ -1,9 +1,10 @@
 model_list:
-  - model_name: azure/any
+  - model_name: fake-openai-endpoint
     litellm_params:
       model: openai/fake
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app/
 
+
 litellm_settings:
   cache: true

litellm/proxy/auth/auth_checks.py

@@ -756,7 +756,9 @@ async def _cache_management_object(
     user_api_key_cache: DualCache,
     proxy_logging_obj: Optional[ProxyLogging],
 ):
-    await user_api_key_cache.async_set_cache(key=key, value=value)
+    await user_api_key_cache.async_set_cache(
+        key=key, value=value, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
+    )
 
 
 async def _cache_team_object(
@@ -870,7 +872,6 @@ async def _get_team_object_from_user_api_key_cache(
         proxy_logging_obj=proxy_logging_obj,
     )
 
-    # save to db access time
     # save to db access time
     _update_last_db_access_time(
         key=db_access_time_key,

tests/litellm/caching/test_in_memory_cache.py

@@ -69,3 +69,22 @@ def test_in_memory_cache_ttl():
     cached_obj = in_memory_cache.get_cache(key="new-fake-key")
     new_ttl_time = in_memory_cache.ttl_dict.get("new-fake-key")
     assert new_ttl_time is None
+
+
+def test_in_memory_cache_ttl_allow_override():
+    """
+    Check that
+    - if ttl is not set, it will be set to default ttl
+    - if object expires, the ttl is also removed
+    """
+    in_memory_cache = InMemoryCache()
+    ## On object expiration, but no get_cache, the override should be allowed
+    in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value", ttl=1)
+    initial_ttl_time = in_memory_cache.ttl_dict["new-fake-key"]
+    assert initial_ttl_time is not None
+    time.sleep(1)
+
+    in_memory_cache.set_cache(key="new-fake-key", value="new-fake-value-2", ttl=1)
+    new_ttl_time = in_memory_cache.ttl_dict["new-fake-key"]
+    assert new_ttl_time is not None
+    assert new_ttl_time != initial_ttl_time