From ec9b84d38c7053be561d387f03b33207482e4eb9 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 22:52:54 +0000 Subject: [PATCH] chore(container,skills): LRU eviction for owner caches; widen file_purpose Literal Two cleanups from the /simplify pass: * ``_CONTAINER_OWNER_CACHE`` and ``_SKILL_CACHE`` now LRU-evict via ``OrderedDict.popitem(last=False)`` instead of full ``clear()`` at capacity. Full clears converted a steady-state cached workload into a periodic full-DB-load oscillation as the cache repopulated from zero and cleared again. Reads now ``move_to_end`` so the just-touched entry survives the next eviction. Mirrors the pre-existing LRU pattern in ``_remember_container_owner``. * ``LiteLLM_ManagedObjectTable.file_purpose`` Literal now includes ``"container"`` so Pydantic validation accepts rows written by the ownership store. --- litellm/llms/litellm_proxy/skills/handler.py | 12 ++++++++--- litellm/proxy/_types.py | 2 +- .../proxy/container_endpoints/ownership.py | 18 ++++++++++++----- .../test_container_proxy_ownership.py | 20 ++++++++++++++++--- .../litellm_proxy/test_skills_ownership.py | 3 ++- 5 files changed, 42 insertions(+), 13 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 7a6dbb175f..e5caf3cffc 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -8,6 +8,7 @@ Used by the transformation layer and skills injection hook. import os import time import uuid +from collections import OrderedDict from typing import Any, Dict, List, Optional, Tuple from litellm._logging import verbose_logger @@ -28,7 +29,7 @@ ALLOW_UNOWNED_SKILL_ACCESS_ENV = "LITELLM_ALLOW_UNOWNED_SKILL_ACCESS" # round-trip per request. Same shape as `_byok_cred_cache` and the container # ownership cache: (value, monotonic_timestamp). `None` is cached as a true # negative ("skill does not exist") so repeated misses also avoid the DB. -_SKILL_CACHE: Dict[str, Tuple[Optional[Any], float]] = {} +_SKILL_CACHE: "OrderedDict[str, Tuple[Optional[Any], float]]" = OrderedDict() _SKILL_CACHE_TTL = 60 # seconds _SKILL_CACHE_MAX_SIZE = 10000 @@ -42,13 +43,18 @@ def _read_skill_cache(skill_id: str) -> Tuple[bool, Optional[Any]]: if time.monotonic() - timestamp > _SKILL_CACHE_TTL: _SKILL_CACHE.pop(skill_id, None) return False, None + _SKILL_CACHE.move_to_end(skill_id) return True, value def _write_skill_cache(skill_id: str, skill: Optional[Any]) -> None: - if len(_SKILL_CACHE) >= _SKILL_CACHE_MAX_SIZE: - _SKILL_CACHE.clear() + # LRU eviction (popitem(last=False)) instead of full ``clear()`` — + # see container ownership cache for rationale. + if skill_id in _SKILL_CACHE: + _SKILL_CACHE.move_to_end(skill_id) _SKILL_CACHE[skill_id] = (skill, time.monotonic()) + while len(_SKILL_CACHE) > _SKILL_CACHE_MAX_SIZE: + _SKILL_CACHE.popitem(last=False) def _invalidate_skill_cache(skill_id: str) -> None: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3cca23f07a..df0a2bcee5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4609,7 +4609,7 @@ class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): unified_object_id: str model_object_id: str - file_purpose: Literal["batch", "fine-tune", "response"] + file_purpose: Literal["batch", "fine-tune", "response", "container"] file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e5577c4052..4905c6537d 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -26,9 +26,10 @@ _IN_MEMORY_CONTAINER_OWNERS: "OrderedDict[str, str]" = OrderedDict() # Short-lived cache keeps every container access check from hitting the DB # (`_get_container_owner` is invoked on retrieve / delete / list / file-content # paths). Mirrors the `_byok_cred_cache` pattern in mcp_server/server.py: -# (value, monotonic_timestamp) tuples, TTL'd, capped, invalidated by writes. -# A `None` value caches "untracked" so repeated negative lookups also avoid DB. -_CONTAINER_OWNER_CACHE: Dict[str, Tuple[Optional[str], float]] = {} +# (value, monotonic_timestamp) tuples, TTL'd, LRU-evicted at capacity, +# invalidated by writes. A ``None`` value caches "untracked" so repeated +# negative lookups also avoid DB. +_CONTAINER_OWNER_CACHE: "OrderedDict[str, Tuple[Optional[str], float]]" = OrderedDict() _CONTAINER_OWNER_CACHE_TTL = 60 # seconds _CONTAINER_OWNER_CACHE_MAX_SIZE = 10000 @@ -42,13 +43,20 @@ def _read_container_owner_cache(model_object_id: str) -> Tuple[bool, Optional[st if time.monotonic() - timestamp > _CONTAINER_OWNER_CACHE_TTL: _CONTAINER_OWNER_CACHE.pop(model_object_id, None) return False, None + _CONTAINER_OWNER_CACHE.move_to_end(model_object_id) return True, value def _write_container_owner_cache(model_object_id: str, owner: Optional[str]) -> None: - if len(_CONTAINER_OWNER_CACHE) >= _CONTAINER_OWNER_CACHE_MAX_SIZE: - _CONTAINER_OWNER_CACHE.clear() + # LRU eviction (popitem(last=False)) instead of full ``clear()`` — a + # full clear at capacity converts a steady-state cached workload into + # a periodic full-DB-load oscillation as the cache repopulates from + # zero and clears again. + if model_object_id in _CONTAINER_OWNER_CACHE: + _CONTAINER_OWNER_CACHE.move_to_end(model_object_id) _CONTAINER_OWNER_CACHE[model_object_id] = (owner, time.monotonic()) + while len(_CONTAINER_OWNER_CACHE) > _CONTAINER_OWNER_CACHE_MAX_SIZE: + _CONTAINER_OWNER_CACHE.popitem(last=False) def _invalidate_container_owner_cache(model_object_id: str) -> None: diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 2a760edf7d..2d175ceedb 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -1228,10 +1228,24 @@ def test_container_owner_cache_expires_after_ttl(monkeypatch): def test_container_owner_cache_evicts_when_at_capacity(monkeypatch): - """The cache must not grow unbounded; reaching capacity clears all entries.""" + """The cache must not grow unbounded; reaching capacity LRU-evicts the + oldest entry, not the entire cache.""" monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) ownership._write_container_owner_cache("a", "user-a") ownership._write_container_owner_cache("b", "user-b") ownership._write_container_owner_cache("c", "user-c") - # Reaching the cap clears everything — the new write is the only survivor. - assert ownership._CONTAINER_OWNER_CACHE.keys() == {"c"} + # ``a`` was the oldest and is dropped; ``b`` and ``c`` survive. + assert list(ownership._CONTAINER_OWNER_CACHE.keys()) == ["b", "c"] + + +def test_container_owner_cache_read_marks_as_recently_used(monkeypatch): + """Reading an entry should reset its position so a subsequent eviction + drops a less-recently-used entry instead of the just-touched one.""" + monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) + ownership._write_container_owner_cache("a", "user-a") + ownership._write_container_owner_cache("b", "user-b") + # Touch ``a`` so it becomes the most-recently-used. + ownership._read_container_owner_cache("a") + ownership._write_container_owner_cache("c", "user-c") + # ``b`` is the LRU at this point; ``a`` and ``c`` survive. + assert list(ownership._CONTAINER_OWNER_CACHE.keys()) == ["a", "c"] diff --git a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py index f98a04205b..d6c0601a57 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -537,4 +537,5 @@ def test_skill_cache_evicts_when_at_capacity(monkeypatch): skills_handler._write_skill_cache("a", Mock()) skills_handler._write_skill_cache("b", Mock()) skills_handler._write_skill_cache("c", Mock()) - assert skills_handler._SKILL_CACHE.keys() == {"c"} + # ``a`` was the oldest and is LRU-evicted; ``b`` and ``c`` survive. + assert list(skills_handler._SKILL_CACHE.keys()) == ["b", "c"]