From a78d121ef40e5e4fa88768f031827ac4aee86b32 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 24 May 2025 20:33:06 -0700 Subject: [PATCH] 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 | 14 +++++++++++++- .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/_new_secret_config.yaml | 3 ++- litellm/proxy/auth/auth_checks.py | 5 +++-- tests/litellm/caching/test_in_memory_cache.py | 19 +++++++++++++++++++ 5 files changed, 37 insertions(+), 5 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/onboarding.html diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 8329e4be54..e9c3f7ba44 100644 --- a/litellm/caching/in_memory_cache.py +++ b/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: diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 6e8b9f7277..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 3c4ecc2ce7..e85ddfdc2b 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/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 \ No newline at end of file diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8653e40c7a..5021247145 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/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, diff --git a/tests/litellm/caching/test_in_memory_cache.py b/tests/litellm/caching/test_in_memory_cache.py index 0dc029d145..72f264b8b7 100644 --- a/tests/litellm/caching/test_in_memory_cache.py +++ b/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