From 9f4c5b75803bcb0201e451ad66e7a19a9f798d7c Mon Sep 17 00:00:00 2001 From: Arthur Moura Carvalho Date: Sun, 31 Aug 2025 12:24:25 -0300 Subject: [PATCH] Fix cache resource leaks and test failures - Add proper cleanup methods to Cache class (close, __del__) - Add context manager support (__enter__, __exit__) for automatic resource management - Update test fixtures to properly cleanup cache instances - Fix SQLite database and file descriptor exhaustion issues during test runs These changes resolve intermittent test failures caused by resource leaks in the FanoutCache implementation, ensuring tests run reliably both individually and as part of the full test suite. --- dspy/clients/cache.py | 29 +++++++++++++++++++++ tests/clients/test_cache.py | 52 +++++++++++++++++++------------------ tests/clients/test_lm.py | 12 ++++++--- 3 files changed, 64 insertions(+), 29 deletions(-) diff --git a/dspy/clients/cache.py b/dspy/clients/cache.py index 187043538e..cca435ae36 100644 --- a/dspy/clients/cache.py +++ b/dspy/clients/cache.py @@ -168,6 +168,35 @@ def load_memory_cache(self, filepath: str) -> None: with open(filepath, "rb") as f: self.memory_cache = cloudpickle.load(f) + def close(self) -> None: + """Close the cache and release all resources.""" + if self.enable_disk_cache and hasattr(self.disk_cache, "close"): + try: + self.disk_cache.close() + except Exception as e: + logger.debug(f"Failed to close disk cache: {e}") + + if self.enable_memory_cache and hasattr(self.memory_cache, "clear"): + with self._lock: + self.memory_cache.clear() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit - ensures cleanup.""" + self.close() + return False + + def __del__(self): + """Destructor to ensure cleanup on garbage collection.""" + try: + self.close() + except Exception: + # Ignore errors during garbage collection + pass + def request_cache( cache_arg_name: str | None = None, diff --git a/tests/clients/test_cache.py b/tests/clients/test_cache.py index 42a1213957..a3a09505ba 100644 --- a/tests/clients/test_cache.py +++ b/tests/clients/test_cache.py @@ -30,44 +30,47 @@ def cache_config(tmp_path): @pytest.fixture def cache(cache_config): """Create a cache instance with the default configuration.""" - return Cache(**cache_config) + cache_instance = Cache(**cache_config) + yield cache_instance + # Cleanup + cache_instance.close() def test_initialization(tmp_path): """Test different cache initialization configurations.""" # Test memory-only cache - memory_cache = Cache( + with Cache( enable_disk_cache=False, enable_memory_cache=True, disk_cache_dir="", disk_size_limit_bytes=0, memory_max_entries=50, - ) - assert isinstance(memory_cache.memory_cache, LRUCache) - assert memory_cache.memory_cache.maxsize == 50 - assert memory_cache.disk_cache == {} + ) as memory_cache: + assert isinstance(memory_cache.memory_cache, LRUCache) + assert memory_cache.memory_cache.maxsize == 50 + assert memory_cache.disk_cache == {} # Test disk-only cache - disk_cache = Cache( + with Cache( enable_disk_cache=True, enable_memory_cache=False, disk_cache_dir=str(tmp_path), disk_size_limit_bytes=1024, memory_max_entries=0, - ) - assert isinstance(disk_cache.disk_cache, FanoutCache) - assert disk_cache.memory_cache == {} + ) as disk_cache: + assert isinstance(disk_cache.disk_cache, FanoutCache) + assert disk_cache.memory_cache == {} # Test disabled cache - disabled_cache = Cache( + with Cache( enable_disk_cache=False, enable_memory_cache=False, disk_cache_dir="", disk_size_limit_bytes=0, memory_max_entries=0, - ) - assert disabled_cache.memory_cache == {} - assert disabled_cache.disk_cache == {} + ) as disabled_cache: + assert disabled_cache.memory_cache == {} + assert disabled_cache.disk_cache == {} def test_cache_key_generation(cache): @@ -180,22 +183,21 @@ def test_save_and_load_memory_cache(cache, tmp_path): cache.save_memory_cache(str(temp_cache_file)) # Create a new cache instance with disk cache disabled - new_cache = Cache( + with Cache( enable_memory_cache=True, enable_disk_cache=False, disk_cache_dir=tmp_path / "disk_cache", disk_size_limit_bytes=0, memory_max_entries=100, - ) - - # Load the memory cache - new_cache.load_memory_cache(str(temp_cache_file)) - - # Verify items are in the new memory cache - for req in requests: - result = new_cache.get(req) - assert result is not None - assert result == f"Response {requests.index(req)}" + ) as new_cache: + # Load the memory cache + new_cache.load_memory_cache(str(temp_cache_file)) + + # Verify items are in the new memory cache + for req in requests: + result = new_cache.get(req) + assert result is not None + assert result == f"Response {requests.index(req)}" def test_request_cache_decorator(cache): diff --git a/tests/clients/test_lm.py b/tests/clients/test_lm.py index 491c3c5940..9d7b586dbf 100644 --- a/tests/clients/test_lm.py +++ b/tests/clients/test_lm.py @@ -24,7 +24,7 @@ def make_response(output_blocks): model="openai/dspy-test-model", object="response", output=output_blocks, - metadata = {}, + metadata={}, parallel_tool_calls=False, temperature=1.0, tool_choice="auto", @@ -92,6 +92,8 @@ def test_dspy_cache(litellm_test_server, tmp_path): assert len(usage_tracker.usage_data) == 0 + # Cleanup + cache.close() dspy.cache = original_cache @@ -213,6 +215,7 @@ def test_reasoning_model_token_parameter(): assert "max_tokens" in lm.kwargs assert lm.kwargs["max_tokens"] == 1000 + @pytest.mark.parametrize("model_name", ["openai/o1", "openai/gpt-5-nano"]) def test_reasoning_model_requirements(model_name): # Should raise assertion error if temperature or max_tokens requirements not met @@ -369,6 +372,8 @@ async def test_async_lm_call_with_cache(tmp_path): assert len(cache.memory_cache) == 2 assert mock_alitellm_completion.call_count == 2 + # Cleanup + cache.close() dspy.cache = original_cache @@ -407,6 +412,7 @@ def test_disable_history(): model="openai/gpt-4o-mini", ) + def test_responses_api(litellm_test_server): api_base, _ = litellm_test_server expected_text = "This is a test answer from responses API." @@ -418,9 +424,7 @@ def test_responses_api(litellm_test_server): "type": "message", "role": "assistant", "status": "completed", - "content": [ - {"type": "output_text", "text": expected_text, "annotations": []} - ], + "content": [{"type": "output_text", "text": expected_text, "annotations": []}], } ] )