diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index b3e6d2eb291..d621d09b439 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -16,7 +16,7 @@ import logging import time import weakref from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence -from contextvars import ContextVar, Token +from contextvars import ContextVar, Token, copy_context from dataclasses import dataclass, field from datetime import timedelta from enum import Enum, auto @@ -528,6 +528,7 @@ class RequestRedisBatches: "_deadline", "_deadline_flush", "_native", + "_native_context", "_native_factory", "_native_flush_hooks", "_native_misses", @@ -551,6 +552,7 @@ class RequestRedisBatches: self._native_factory: Final = native_factory self._native_send_factory: Final = native_send_factory self._native: NativeRequestRedis | None = None + self._native_context: Final = copy_context() self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry self._native_misses: frozenset[str] = frozenset() self._deadline: asyncio.TimerHandle | None = None @@ -580,9 +582,12 @@ class RequestRedisBatches: return batch def native(self) -> NativeRequestRedis: - """The request's native batches, built on first use so requests that touch no native store pay nothing.""" + """The request's native batches, built on first use so requests that touch no native store pay nothing. + + They keep a copy of the context they are built in, and the cycle collector cannot see through native + objects, so they are built in the context this scope was opened from, which cannot reach the scope.""" if self._native is None: - self._native = self._native_factory(self.post_call_deadline) + self._native = self._native_context.run(self._native_factory, self.post_call_deadline) return self._native def native_post_call(self) -> NativeRequestRedis: diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index 536d510f6d6..158b9754309 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -3,9 +3,11 @@ from __future__ import annotations import asyncio +import gc import hashlib import json import uuid +import weakref from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass, field from datetime import timedelta @@ -19,6 +21,7 @@ from litellm._service_logger import ServiceLogging from litellm.caching.redis_batch import ( MIXED_PIPELINE_TARGET, RedisBatch, + RequestRedisBatches, active_request_redis_batch, request_redis_batch_scope, ) @@ -418,6 +421,17 @@ def test_request_scope_hands_out_one_batch_per_backend_and_nests() -> None: assert active_request_redis_batch(cache_a) is None +def test_a_request_scope_that_built_its_native_batches_is_freed_once_it_closes() -> None: + def serve_request() -> weakref.ref[RequestRedisBatches]: + with request_redis_batch_scope() as batches: + batches.native() + return weakref.ref(batches) + + closed: Final = serve_request() + gc.collect() + assert closed() is None + + @pytest.mark.asyncio async def test_a_key_an_mget_read_as_absent_stays_known_missing_until_something_sets_it() -> None: cache, client = make()