From 3e6c55da51096946b783d557df3aa7700aeb12d2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 19:49:22 +0000 Subject: [PATCH] refactor(redis): isolate managed object metadata store --- .../proxy/hooks/managed_files.py | 25 ++- litellm/redis/stores/managed_objects.py | 32 +++ .../proxy/test_managed_files_hook.py | 11 +- .../unit/redis/stores/test_managed_objects.py | 209 ++++++++++++++++++ 4 files changed, 268 insertions(+), 9 deletions(-) create mode 100644 litellm/redis/stores/managed_objects.py create mode 100644 tests/unit/redis/stores/test_managed_objects.py diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index f5ac8718be6..d7f075bc9ae 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -95,6 +95,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_metadata, ) from litellm.redis.stores.managed_files import ManagedFileStore, TransportManagedFileStore, file_key +from litellm.redis.stores.managed_objects import ManagedObjectStore, TransportManagedObjectStore, object_key from litellm.redis.transport.contract import RedisFailure from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.responses.utils import ResponsesAPIRequestUtils @@ -655,6 +656,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await local_cache.async_set_cache(file_key(file_id), value, local_only=True) return value + async def _cache_object(self, unified_object_id: str, record: Mapping[str, object]) -> None: + cache: Final = self._metadata_cache() + local_cache: Final = cache.local + await local_cache.async_set_cache(object_key(unified_object_id), record, local_only=True) + if cache.redis is None: + return + store: Final[ManagedObjectStore] = TransportManagedObjectStore(cache.redis.transport()) + try: + seconds: Final = ( + int(cache.in_memory_ttl) + if cache.in_memory_ttl is not None and not math.isnan(cache.in_memory_ttl) + else cache.redis.default_ttl + ) + ttl: Final = timedelta(seconds=seconds) + except (OverflowError, ValueError): + return + await store.write_object(unified_object_id, record, ttl=ttl) + @staticmethod def _get_prometheus_logger(): """Find PrometheusLogger from litellm.callbacks, if registered.""" @@ -785,11 +804,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose=file_purpose, file_object=stored_object, ) - await self.internal_usage_cache.async_set_cache( - key=unified_object_id, - value=litellm_managed_object.model_dump(), - litellm_parent_otel_span=litellm_parent_otel_span, - ) + await self._cache_object(unified_object_id, litellm_managed_object.model_dump()) from prisma import Json diff --git a/litellm/redis/stores/managed_objects.py b/litellm/redis/stores/managed_objects.py new file mode 100644 index 00000000000..f1b4c2d7579 --- /dev/null +++ b/litellm/redis/stores/managed_objects.py @@ -0,0 +1,32 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +def object_key(unified_object_id: str) -> str: + return unified_object_id + + +class ManagedObjectStore(Protocol): + async def write_object( + self, unified_object_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportManagedObjectStore: + transport: RedisTransport + + async def write_object( + self, unified_object_id: str, record: Mapping[str, object], *, ttl: timedelta | None + ) -> None | RedisFailure: + try: + encoded: Final = json.dumps(record) + except (TypeError, ValueError) as exc: + return RedisFailure(kind="decode", message=str(exc)) + result: Final = await self.transport.set_seconds(object_key(unified_object_id), encoded, ttl=ttl) + return result if isinstance(result, RedisFailure) else None diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index c0f6018469f..055583322a1 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1380,8 +1380,9 @@ def _make_object_store_instance(): PROXY_LiteLLMManagedFiles, ) - mock_cache = MagicMock() - mock_cache.async_set_cache = AsyncMock() + from litellm.caching.dual_cache import DualCache + + cache = DualCache() mock_prisma = MagicMock() mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() @@ -1389,7 +1390,7 @@ def _make_object_store_instance(): return ( PROXY_LiteLLMManagedFiles( - internal_usage_cache=mock_cache, + internal_usage_cache=cache, prisma_client=mock_prisma, ), mock_prisma, @@ -1498,7 +1499,9 @@ async def test_a_stored_background_response_names_its_creators_end_user_and_hold ) written = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"] - cached = managed_files.internal_usage_cache.async_set_cache.await_args.kwargs["value"] + cached: object = await managed_files.internal_usage_cache.async_get_cache( + "resp_unified", litellm_parent_otel_span=None + ) assert generated_text not in json.dumps(written, default=str) assert generated_text not in json.dumps(cached, default=str) assert json.loads(written["create"]["file_object"])["status"] == "completed" diff --git a/tests/unit/redis/stores/test_managed_objects.py b/tests/unit/redis/stores/test_managed_objects.py new file mode 100644 index 00000000000..05cd5e5d7b2 --- /dev/null +++ b/tests/unit/redis/stores/test_managed_objects.py @@ -0,0 +1,209 @@ +import json +import time +from collections.abc import Mapping +from datetime import timedelta +from types import SimpleNamespace +from typing import Final, Literal, cast +from unittest.mock import patch + +import fakeredis +import pytest +from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.models.managed_files import LiteLLM_ManagedObjectTable +from litellm.proxy.utils import InternalUsageCache, PrismaClient +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.utils import LiteLLMBatch, LiteLLMFineTuningJob +from tests._support.fake_redis import fake_redis_cache + + +class _Table: + def __init__(self, raw: fakeredis.FakeAsyncRedis | None = None, expected: bytes | None = None) -> None: + self.raw = raw + self.expected = expected + self.created: Mapping[str, object] | None = None + self.updated: Mapping[str, object] | None = None + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> None: + if self.raw is not None: + assert await self.raw.get(f"tenant:{where['unified_object_id']}") == self.expected + self.created = data + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: + if self.raw is not None: + assert await self.raw.get(f"tenant:{where['unified_object_id']}") == self.expected + self.updated = data + return 0 + + +def _hook(cache: DualCache, table: _Table) -> PROXY_LiteLLMManagedFiles: + db: Final = SimpleNamespace(litellm_managedobjecttable=table) + return PROXY_LiteLLMManagedFiles( + InternalUsageCache(cache), cast(PrismaClient, SimpleNamespace(db=db, writer_db=db)) + ) + + +def _object( + purpose: Literal["batch", "fine-tune", "response"], +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + if purpose == "batch": + return LiteLLMBatch( + id="batch-é", + completion_window="24h", + created_at=123, + endpoint="/v1/chat/completions", + input_file_id="input-文", + object="batch", + status="completed", + ) + if purpose == "fine-tune": + return LiteLLMFineTuningJob( + id="ft-é", + object="fine_tuning.job", + model="model", + created_at=123, + organization_id="org", + result_files=[], + status="validating_files", + training_file="file-文", + hyperparameters={"n_epochs": 8}, + seed=0, + ) + return ResponsesAPIResponse( + id="resp-é", + created_at=123, + status="completed", + output=[ + { + "type": "message", + "id": "msg", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "private generated output", "annotations": []}], + } + ], + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("purpose", ["batch", "fine-tune", "response"]) +@pytest.mark.parametrize("ttl", [1, 21.9, None, 0.5, 0, -1]) +async def test_object_wire_and_ttl_match_old_reader_and_write_precedes_database( + server: fakeredis.FakeServer, + raw: fakeredis.FakeAsyncRedis, + purpose: Literal["batch", "fine-tune", "response"], + ttl: float | None, +) -> None: + value: Final = _object(purpose) + sanitized: Final = value.model_copy(update={"output": []}) if isinstance(value, ResponsesAPIResponse) else value + record: Final = LiteLLM_ManagedObjectTable( + unified_object_id="object-é", + model_object_id=value.id, + file_purpose=purpose, + file_object=sanitized, + ).model_dump() + expected: Final = None if ttl is not None and ttl < 1 else json.dumps(record).encode() + cache: Final = fake_redis_cache(server, namespace="tenant") + table: Final = _Table(raw, expected) + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=ttl, default_redis_ttl=999) + dual.update_cache_ttl(ttl, None) + hook: Final = _hook(dual, table) + with patch("time.time", return_value=time.time()): + await hook.store_unified_object_id("object-é", value, None, value.id, purpose, UserAPIKeyAuth(user_id="owner")) + await cache.async_set_cache("legacy", record, **({"ttl": ttl} if ttl is not None else {})) + assert await raw.get("tenant:object-é") == await raw.get("tenant:legacy") == expected + assert await raw.pttl("tenant:object-é") == await raw.pttl("tenant:legacy") + assert await cache.async_get_cache("object-é") == (None if ttl is not None and ttl < 1 else record) + assert await raw.get("tenant:tenant:object-é") is None + assert await dual.async_get_cache("object-é", local_only=True) == ( + None if ttl is not None and ttl < 0 else record + ) + assert table.created is not None + assert "private generated output" not in json.dumps(table.created, default=str) + if isinstance(value, ResponsesAPIResponse): + assert len(value.output) == 1 # pyright: ignore[reportUnknownArgumentType] # output carries a bare dict member + + +@pytest.mark.asyncio +async def test_object_update_only_still_writes_and_expires_cache( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + value: Final = _object("batch") + with patch("time.time", return_value=1_800_000_000): + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + table: Final = _Table() + hook: Final = _hook(dual, table) + await hook.store_unified_object_id( + "object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="poller"), create_if_missing=False + ) + assert await raw.pttl("object") == 1000 + assert table.created is None + assert table.updated is not None + assert "created_by" not in table.updated + with patch("time.time", return_value=1_800_000_002): + assert await raw.get("object") is None + assert await dual.async_get_cache("object", local_only=True) is None + + +@pytest.mark.asyncio +async def test_object_outage_keeps_local_and_database_write(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache) + table: Final = _Table() + hook: Final = _hook(dual, table) + value: Final = _object("batch") + server.connected = False + for _ in range(8): + await hook.store_unified_object_id("object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="owner")) + assert table.created is not None + assert ( + await dual.async_get_cache("object", local_only=True) + == LiteLLM_ManagedObjectTable( + unified_object_id="object", + model_object_id=value.id, + file_purpose="batch", + file_object=value, + ).model_dump() + ) + assert cache._circuit_breaker.state == "open" # pyright: ignore[reportPrivateUsage] # asserts the breaker opened after the outage + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ttl", [float("nan"), float("inf"), float("-inf")]) +async def test_invalid_object_cache_ttl_does_not_block_database_write( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, ttl: float +) -> None: + table: Final = _Table() + dual: Final = DualCache(redis_cache=fake_redis_cache(server), default_in_memory_ttl=ttl) + hook: Final = _hook(dual, table) + value: Final = _object("batch") + with patch("time.time", return_value=1_800_000_000): + await dual.async_set_cache("legacy", value.model_dump()) + await hook.store_unified_object_id("object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="owner")) + assert table.created is not None + assert (await raw.get("object") is None) == (await raw.get("legacy") is None) + assert await raw.pttl("object") == await raw.pttl("legacy") + + +@pytest.mark.asyncio +async def test_object_unknown_applied_write_is_not_undone( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + cache: Final = fake_redis_cache(server) + dual: Final = DualCache(redis_cache=cache, default_in_memory_ttl=1) + table: Final = _Table() + hook: Final = _hook(dual, table) + value: Final = _object("batch") + + async def applied_then_timeout(name: str, value: str, *, ex: timedelta | None, nx: bool) -> bool: + await raw.set(name, value, ex=ex, nx=nx) + raise RedisTimeoutError("reply lost after apply") + + with patch.object(cache.init_async_client(), "set", side_effect=applied_then_timeout): + await hook.store_unified_object_id("object", value, None, value.id, "batch", UserAPIKeyAuth(user_id="owner")) + assert await raw.get("object") == json.dumps(await dual.async_get_cache("object", local_only=True)).encode() + assert table.created is not None