# SPDX-License-Identifier: Apache-2.0 """Tests for the in-memory hot cache tier in PagedSSDCacheManager.""" import threading import time from pathlib import Path from typing import List from unittest.mock import patch import pytest from omlx.cache.paged_ssd_cache import ( PagedSSDBlockMetadata, PagedSSDCacheManager, SharedHotCacheBudget, _extract_tensor_bytes, ) try: import mlx.core as mx HAS_MLX = True except ImportError: HAS_MLX = False @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCacheDisabled: """Verify that hot_cache_max_bytes=0 preserves existing behaviour.""" @pytest.fixture def manager(self, tmp_path): mgr = PagedSSDCacheManager( cache_dir=tmp_path / "ssd_cache", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=0, ) yield mgr mgr.close() def test_hot_cache_disabled_by_default(self, manager): """hot_cache_max_bytes=0 means hot cache is disabled.""" assert manager._hot_cache_enabled is False assert manager._hot_cache_max_bytes == 0 def test_save_load_works_without_hot_cache(self, manager): """Save/load should work even when hot cache is disabled.""" block_hash = b"disabled_hot_cache_test" cache_data = [ (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64))) for _ in range(4) ] result = manager.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name="test-model", layer_cache_types=["KVCache"] * 4, ) assert result is True assert manager.has_block(block_hash) loaded = manager.load_block(block_hash) assert loaded is not None assert len(loaded) == 4 def test_stats_hot_cache_zero_when_disabled(self, manager): """Hot cache stats should be zero when disabled.""" stats = manager.get_stats() assert stats.hot_cache_entries == 0 assert stats.hot_cache_size_bytes == 0 assert stats.hot_cache_max_bytes == 0 assert stats.hot_cache_hits == 0 assert stats.hot_cache_evictions == 0 assert stats.hot_cache_promotions == 0 @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCacheEnabled: """Test hot cache with in-memory caching active.""" @pytest.fixture def manager(self, tmp_path): # 10 MB hot cache — generous for test blocks mgr = PagedSSDCacheManager( cache_dir=tmp_path / "ssd_cache", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) yield mgr mgr.close() def _make_cache_data(self, num_layers=4, seq_len=32, heads=4, head_dim=32): """Create test cache data.""" return [ ( mx.zeros((1, heads, seq_len, head_dim)), mx.zeros((1, heads, seq_len, head_dim)), ) for _ in range(num_layers) ] def _save_block(self, manager, block_hash, num_layers=4, model="test-model"): """Save a test block and return True on success.""" cache_data = self._make_cache_data(num_layers=num_layers) return manager.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name=model, layer_cache_types=["KVCache"] * num_layers, ) def test_save_stores_in_hot_cache(self, manager): """After save_block(), the entry should be in hot cache.""" block_hash = b"hot_cache_save_test1" self._save_block(manager, block_hash) # Verify hot cache has the entry entry = manager._hot_cache_get(block_hash) assert entry is not None assert "tensors_raw" in entry assert entry["num_layers"] == 4 def test_load_from_hot_cache(self, manager): """load_block() should return data from hot cache without SSD I/O.""" block_hash = b"hot_cache_load_test1" self._save_block(manager, block_hash) loaded = manager.load_block(block_hash) assert loaded is not None assert len(loaded) == 4 stats = manager.get_stats() assert stats.hot_cache_hits >= 1 def test_hot_cache_hit_updates_stats(self, manager): """Hot cache hit should increment hot_cache_hits counter.""" block_hash = b"hot_cache_stats_test1" self._save_block(manager, block_hash) initial_stats = manager.get_stats() initial_hits = initial_stats.hot_cache_hits manager.load_block(block_hash) manager.load_block(block_hash) stats = manager.get_stats() assert stats.hot_cache_hits >= initial_hits + 2 def test_hot_cache_size_tracking(self, manager): """Hot cache should track total size in bytes.""" block_hash = b"hot_cache_size_test1" self._save_block(manager, block_hash) stats = manager.get_stats() assert stats.hot_cache_entries == 1 assert stats.hot_cache_size_bytes > 0 assert stats.hot_cache_max_bytes == 10 * 1024**2 def test_delete_block_removes_from_hot_cache(self, manager): """delete_block() should remove entry from hot cache.""" block_hash = b"hot_cache_delete_test" self._save_block(manager, block_hash) # Verify it's in hot cache assert manager._hot_cache_get(block_hash) is not None manager.delete_block(block_hash) # Verify it's gone from hot cache assert manager._hot_cache_get(block_hash) is None def test_close_clears_hot_cache(self, tmp_path): """close() should clear all hot cache entries.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "close_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) block_hash = b"hot_cache_close_test1" cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name="test", layer_cache_types=["KVCache"] * 4, ) assert len(mgr._hot_cache) > 0 mgr.close() assert len(mgr._hot_cache) == 0 assert mgr._hot_cache_total_bytes == 0 @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCacheLRU: """Test LRU eviction behaviour of the hot cache.""" def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16): """Create small test cache data.""" return [ ( mx.zeros((1, heads, seq_len, head_dim)), mx.zeros((1, heads, seq_len, head_dim)), ) for _ in range(num_layers) ] def _entry_size(self, num_layers=2, seq_len=16, heads=2, head_dim=16): """Estimate the raw byte size of one entry.""" # Each tensor: 1 * heads * seq_len * head_dim * 4 bytes (float32) # 2 tensors (keys + values) per layer, num_layers layers return num_layers * 2 * 1 * heads * seq_len * head_dim * 4 def test_lru_eviction(self, tmp_path): """Old entries should be evicted when capacity is exceeded.""" entry_size = self._entry_size() # Allow room for exactly 2 entries max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "lru_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: # Save 3 blocks — the first should be evicted for i in range(3): block_hash = f"lru_block_{i}".encode() cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # Block 0 should have been evicted (LRU) assert mgr._hot_cache_get(b"lru_block_0") is None # Blocks 1 and 2 should still be in hot cache assert mgr._hot_cache_get(b"lru_block_1") is not None assert mgr._hot_cache_get(b"lru_block_2") is not None stats = mgr.get_stats() assert stats.hot_cache_evictions >= 1 finally: mgr.close() def test_lru_access_refreshes_order(self, tmp_path): """Accessing a block should move it to MRU position.""" entry_size = self._entry_size() max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "lru_order_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: # Save blocks 0 and 1 for i in range(2): block_hash = f"order_block_{i}".encode() cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # Access block 0 to refresh its LRU position mgr.load_block(b"order_block_0") # Save block 2 — should evict block 1 (LRU), not block 0 cache_data = self._make_cache_data() mgr.save_block( block_hash=b"order_block_2", cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # Block 0 was accessed so should still be present assert mgr._hot_cache_get(b"order_block_0") is not None # Block 1 was LRU and should be evicted assert mgr._hot_cache_get(b"order_block_1") is None # Block 2 was just added assert mgr._hot_cache_get(b"order_block_2") is not None finally: mgr.close() @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCachePromotion: """Test promotion from SSD to hot cache on load.""" def _make_cache_data(self, num_layers=4, seq_len=32, heads=4, head_dim=32): return [ ( mx.zeros((1, heads, seq_len, head_dim)), mx.zeros((1, heads, seq_len, head_dim)), ) for _ in range(num_layers) ] def test_ssd_load_promotes_to_hot_cache(self, tmp_path): """Loading a block from SSD should promote it to hot cache.""" # Use hot_cache disabled to write directly to SSD first mgr_cold = PagedSSDCacheManager( cache_dir=tmp_path / "promote_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=0, ) block_hash = b"promote_test_block1" cache_data = self._make_cache_data() mgr_cold.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name="test", layer_cache_types=["KVCache"] * 4, ) # Wait for background SSD write to complete time.sleep(0.5) mgr_cold.close() # Now open with hot cache enabled — block is on SSD only mgr = PagedSSDCacheManager( cache_dir=tmp_path / "promote_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) try: assert mgr._hot_cache_get(block_hash) is None # Load from SSD — should promote to hot cache loaded = mgr.load_block(block_hash) assert loaded is not None assert len(loaded) == 4 # Verify promotion happened assert mgr._hot_cache_get(block_hash) is not None stats = mgr.get_stats() assert stats.hot_cache_promotions >= 1 finally: mgr.close() def test_ssd_load_can_skip_hot_cache_promotion(self, tmp_path): """SSD loads can reconstruct a request without retaining a hot RAM copy.""" mgr_cold = PagedSSDCacheManager( cache_dir=tmp_path / "skip_promote_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=0, ) block_hash = b"skip_promote_blk_1" cache_data = self._make_cache_data() mgr_cold.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name="test", layer_cache_types=["KVCache"] * 4, ) mgr_cold.close() mgr = PagedSSDCacheManager( cache_dir=tmp_path / "skip_promote_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) try: loaded, metadata = mgr.load_block_with_metadata( block_hash, promote_to_hot_cache=False, ) assert loaded is not None assert metadata is not None assert mgr._hot_cache_get(block_hash) is None loaded_again, _ = mgr.load_block_with_metadata(block_hash) assert loaded_again is not None assert mgr._hot_cache_get(block_hash) is not None finally: mgr.close() def test_write_through_save_bypasses_hot_cache_retention(self, tmp_path): """Pressure write-through keeps SSD durability without hot-cache RAM.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "write_through_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) try: block_hash = b"write_through_blk1" cache_data = self._make_cache_data() result = mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name="test", layer_cache_types=["KVCache"] * 4, hot_cache_write_back=False, ) assert result is True assert mgr.has_block(block_hash) assert mgr._hot_cache_get(block_hash) is None loaded, metadata = mgr.load_block_with_metadata( block_hash, promote_to_hot_cache=False, ) assert loaded is not None assert metadata is not None assert mgr._hot_cache_get(block_hash) is None finally: mgr.close() def test_promotion_does_not_happen_when_disabled(self, tmp_path): """No promotion when hot cache is disabled.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "no_promote_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=0, ) try: block_hash = b"no_promote_block1__" cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=64, model_name="test", layer_cache_types=["KVCache"] * 4, ) # Wait for background write time.sleep(0.5) # Clear the temporary buffer (simulates what happens after write completes) mgr._hot_cache_remove(block_hash) # Load from SSD loaded = mgr.load_block(block_hash) assert loaded is not None stats = mgr.get_stats() assert stats.hot_cache_promotions == 0 finally: mgr.close() @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCacheCacheTypes: """Test hot cache with various cache types (KVCache, CacheList).""" @pytest.fixture def manager(self, tmp_path): mgr = PagedSSDCacheManager( cache_dir=tmp_path / "types_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) yield mgr mgr.close() def test_cache_list_blocks(self, manager): """Hot cache should handle CacheList blocks correctly.""" block_hash = b"cache_list_hot_test" sub_keys1 = mx.zeros((1, 8, 32, 64)) sub_values1 = mx.ones((1, 8, 32, 64)) sub_keys2 = mx.zeros((1, 4, 32, 64)) sub_values2 = mx.ones((1, 4, 32, 64)) cache_data = [ ( "__cache_list__", [(sub_keys1, sub_values1), (sub_keys2, sub_values2)], ), (mx.zeros((1, 8, 32, 64)), mx.ones((1, 8, 32, 64))), ] result = manager.save_block( block_hash=block_hash, cache_data=cache_data, token_count=32, model_name="test", layer_cache_types=["CacheList", "KVCache"], ) assert result is True # Load from hot cache loaded = manager.load_block(block_hash) assert loaded is not None assert len(loaded) == 2 # First layer is CacheList assert isinstance(loaded[0], list) assert len(loaded[0]) == 2 # Second layer is KVCache tuple assert isinstance(loaded[1], tuple) class TestHotCacheConcurrency: """Test thread safety of hot cache internal operations.""" def test_concurrent_put_get(self, tmp_path): """Hot cache put/get should be thread-safe under concurrent access.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "concurrent_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=50 * 1024**2, ) errors: List[Exception] = [] num_threads = 8 ops_per_thread = 20 def worker(thread_id): try: for i in range(ops_per_thread): block_hash = f"conc_{thread_id}_{i}____".encode() # Create a fake hot cache entry with raw bytes raw_data = bytes(1024) # 1KB of zeros entry = { "tensors_raw": { "layer_0_keys": (raw_data, "float32", [1, 2, 16, 8]), "layer_0_values": (raw_data, "float32", [1, 2, 16, 8]), }, "file_metadata": {}, "num_layers": 1, "layer_cache_types": ["KVCache"], "block_metadata": None, } mgr._hot_cache_put(block_hash, entry) # Read back for i in range(ops_per_thread): block_hash = f"conc_{thread_id}_{i}____".encode() result = mgr._hot_cache_get(block_hash) # May be None if evicted by another thread, that's OK except Exception as e: errors.append(e) threads = [ threading.Thread(target=worker, args=(t,)) for t in range(num_threads) ] for t in threads: t.start() for t in threads: t.join(timeout=30) mgr.close() assert len(errors) == 0, f"Concurrent errors: {errors}" class TestHotCacheByteAccounting: """Regression tests for raw-byte hot cache size accounting.""" def _make_raw_entry( self, tmp_path: Path, block_hash: bytes, model_name: str = "test-model", key_size: int = 128, value_size: int = 256, ): metadata = PagedSSDBlockMetadata( block_hash=block_hash, file_path=tmp_path / f"{block_hash.hex()}.safetensors", file_size=key_size + value_size, token_count=16, created_at=time.time(), last_access=time.time(), num_layers=1, model_name=model_name, layer_cache_types=["KVCache"], ) entry = { "tensors_raw": { "layer_0_keys": (bytes(key_size), "float32", [1, 1, 1, key_size // 4]), "layer_0_values": ( bytes(value_size), "float32", [1, 1, 1, value_size // 4], ), }, "file_metadata": {}, "num_layers": 1, "layer_cache_types": ["KVCache"], "block_metadata": metadata, } return entry, key_size + value_size def test_hot_cache_remove_decrements_raw_entry_bytes(self, tmp_path): """Removing a tensors_raw entry should subtract its bytes from the counter.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "remove_accounting", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=1024**2, hot_cache_only=True, ) try: block_hash = b"remove_bytes_test" entry, expected_size = self._make_raw_entry(tmp_path, block_hash) mgr._hot_cache_put(block_hash, entry) assert mgr._hot_cache_total_bytes == expected_size mgr._hot_cache_remove(block_hash) assert mgr._hot_cache_get(block_hash) is None assert mgr._hot_cache_total_bytes == 0 assert mgr.get_stats().hot_cache_size_bytes == 0 finally: mgr.close() def test_get_stats_for_model_reports_raw_hot_cache_bytes(self, tmp_path): """Per-model stats should count tensors_raw entries in the hot cache.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "model_accounting", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=1024**2, hot_cache_only=True, ) try: model_a_hash = b"model_a_hot_bytes" model_b_hash = b"model_b_hot_bytes" model_a_entry, model_a_size = self._make_raw_entry( tmp_path, model_a_hash, model_name="model-a" ) model_b_entry, _ = self._make_raw_entry( tmp_path, model_b_hash, model_name="model-b" ) mgr._hot_cache_put(model_a_hash, model_a_entry) mgr._hot_cache_put(model_b_hash, model_b_entry) stats = mgr.get_stats_for_model("model-a") assert stats.hot_cache_entries == 1 assert stats.hot_cache_size_bytes == model_a_size finally: mgr.close() def test_shared_budget_evicts_across_managers_by_global_lru(self, tmp_path): """A shared hot cache budget should cap aggregate manager memory.""" entry, entry_size = self._make_raw_entry(tmp_path, b"budget_size_probe") del entry budget = SharedHotCacheBudget(entry_size * 2 + 1) mgr_a = PagedSSDCacheManager( cache_dir=tmp_path / "budget_a", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=budget.max_bytes, hot_cache_only=True, hot_cache_budget=budget, ) mgr_b = PagedSSDCacheManager( cache_dir=tmp_path / "budget_b", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=budget.max_bytes, hot_cache_only=True, hot_cache_budget=budget, ) try: a0, _ = self._make_raw_entry(tmp_path, b"budget_a_0", model_name="model-a") b0, _ = self._make_raw_entry(tmp_path, b"budget_b_0", model_name="model-b") b1, _ = self._make_raw_entry(tmp_path, b"budget_b_1", model_name="model-b") mgr_a._hot_cache_put(b"budget_a_0", a0) mgr_b._hot_cache_put(b"budget_b_0", b0) assert budget.total_bytes == entry_size * 2 # Refresh A, so B's first block is the global LRU victim. assert mgr_a._hot_cache_get(b"budget_a_0") is not None mgr_b._hot_cache_put(b"budget_b_1", b1) assert budget.total_bytes == entry_size * 2 assert mgr_a._hot_cache_get(b"budget_a_0") is not None assert mgr_b._hot_cache_get(b"budget_b_0") is None assert mgr_b._hot_cache_get(b"budget_b_1") is not None assert mgr_a.get_stats().hot_cache_max_bytes == budget.max_bytes assert mgr_b.get_stats().hot_cache_max_bytes == budget.max_bytes finally: mgr_a.close() mgr_b.close() def test_shared_budget_forgets_owner_on_clear_and_close(self, tmp_path): """Clearing or closing a manager must release shared budget bytes.""" _, entry_size = self._make_raw_entry(tmp_path, b"budget_clear_probe") budget = SharedHotCacheBudget(entry_size * 4) mgr_a = PagedSSDCacheManager( cache_dir=tmp_path / "budget_clear_a", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=budget.max_bytes, hot_cache_only=True, hot_cache_budget=budget, ) mgr_b = PagedSSDCacheManager( cache_dir=tmp_path / "budget_clear_b", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=budget.max_bytes, hot_cache_only=True, hot_cache_budget=budget, ) try: a0, _ = self._make_raw_entry(tmp_path, b"budget_clear_a_0") b0, _ = self._make_raw_entry(tmp_path, b"budget_clear_b_0") mgr_a._hot_cache_put(b"budget_clear_a_0", a0) mgr_b._hot_cache_put(b"budget_clear_b_0", b0) assert budget.total_bytes == entry_size * 2 assert mgr_a.clear_hot_cache() == 1 assert budget.total_bytes == entry_size mgr_b.close() assert budget.total_bytes == 0 finally: mgr_a.close() mgr_b.close() def test_shared_budget_shrink_respects_lru_and_protected_hashes(self, tmp_path): """Pressure shrink should skip protected active-request chains.""" _, entry_size = self._make_raw_entry(tmp_path, b"budget_shrink_probe") budget = SharedHotCacheBudget(entry_size * 3) mgr_a = PagedSSDCacheManager( cache_dir=tmp_path / "budget_shrink_a", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=budget.max_bytes, hot_cache_only=True, hot_cache_budget=budget, ) mgr_b = PagedSSDCacheManager( cache_dir=tmp_path / "budget_shrink_b", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=budget.max_bytes, hot_cache_only=True, hot_cache_budget=budget, ) try: a0, _ = self._make_raw_entry(tmp_path, b"budget_shrink_a0") b0, _ = self._make_raw_entry(tmp_path, b"budget_shrink_b0") b1, _ = self._make_raw_entry(tmp_path, b"budget_shrink_b1") mgr_a._hot_cache_put(b"budget_shrink_a0", a0) mgr_b._hot_cache_put(b"budget_shrink_b0", b0) mgr_b._hot_cache_put(b"budget_shrink_b1", b1) freed = budget.shrink_to( entry_size, protected_hashes={b"budget_shrink_b0"}, ) assert freed == entry_size * 2 assert budget.total_bytes == entry_size assert mgr_a._hot_cache_get(b"budget_shrink_a0") is None assert mgr_b._hot_cache_get(b"budget_shrink_b0") is not None assert mgr_b._hot_cache_get(b"budget_shrink_b1") is None finally: mgr_a.close() mgr_b.close() def test_local_hot_cache_shrink_respects_protected_hashes(self, tmp_path): _, entry_size = self._make_raw_entry(tmp_path, b"local_shrink_probe") mgr = PagedSSDCacheManager( cache_dir=tmp_path / "local_shrink", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=entry_size * 3, hot_cache_only=True, ) try: h0 = b"local_shrink_0" h1 = b"local_shrink_1" h2 = b"local_shrink_2" e0, _ = self._make_raw_entry(tmp_path, h0) e1, _ = self._make_raw_entry(tmp_path, h1) e2, _ = self._make_raw_entry(tmp_path, h2) mgr._hot_cache_put(h0, e0) mgr._hot_cache_put(h1, e1) mgr._hot_cache_put(h2, e2) freed = mgr.shrink_hot_cache_to(entry_size, protected_hashes={h1}) assert freed == entry_size * 2 assert mgr.get_stats().hot_cache_size_bytes == entry_size assert mgr._hot_cache_get(h0) is None assert mgr._hot_cache_get(h1) is not None assert mgr._hot_cache_get(h2) is None finally: mgr.close() @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCacheStatsAccuracy: """Test that hot cache statistics are accurate.""" def test_all_stats_counters(self, tmp_path): """Verify hot cache stats counters are correctly maintained.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 # ~4096 bytes per entry # Room for 2 entries max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "stats_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: def save(idx): block_hash = f"stats_block_{idx}__".encode() cache_data = [ (mx.zeros((1, 2, 16, 16)), mx.zeros((1, 2, 16, 16))) for _ in range(2) ] mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # Save 2 blocks: fits in hot cache save(0) save(1) stats = mgr.get_stats() assert stats.hot_cache_entries == 2 assert stats.hot_cache_evictions == 0 # Save 3rd block: triggers eviction of block 0 save(2) stats = mgr.get_stats() assert stats.hot_cache_entries == 2 assert stats.hot_cache_evictions >= 1 # Load block 1 (hot cache hit) mgr.load_block(b"stats_block_1__") stats = mgr.get_stats() assert stats.hot_cache_hits >= 1 # Verify size tracking is positive assert stats.hot_cache_size_bytes > 0 assert stats.hot_cache_max_bytes == max_bytes finally: mgr.close() @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestHotCacheWriteBack: """Test write-back behavior: no SSD writes until eviction or shutdown.""" def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16): return [ ( mx.zeros((1, heads, seq_len, head_dim)), mx.zeros((1, heads, seq_len, head_dim)), ) for _ in range(num_layers) ] def test_save_does_not_write_to_ssd(self, tmp_path): """With hot cache enabled, save_block should not create SSD files.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "wb_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) try: block_hash = b"wb_no_ssd_write_t1" cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # Block should be in hot cache assert mgr._hot_cache_get(block_hash) is not None # No SSD file should exist yet time.sleep(0.3) ssd_files = list((tmp_path / "wb_test").rglob("*.safetensors")) assert len(ssd_files) == 0, f"Unexpected SSD files: {ssd_files}" finally: mgr.close() def test_eviction_writes_to_ssd(self, tmp_path): """When hot cache evicts, the evicted block should be written to SSD.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "wb_evict_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: # Save 2 blocks (fits in hot cache) for i in range(2): block_hash = f"wb_evict_blk_{i}__".encode() cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # No SSD files yet time.sleep(0.3) ssd_files = list((tmp_path / "wb_evict_test").rglob("*.safetensors")) assert len(ssd_files) == 0 # Save 3rd block → evicts block 0 → should trigger SSD write cache_data = self._make_cache_data() mgr.save_block( block_hash=b"wb_evict_blk_2__", cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # Wait for background writer to process the evicted block time.sleep(0.5) ssd_files = list((tmp_path / "wb_evict_test").rglob("*.safetensors")) assert len(ssd_files) >= 1, "Evicted block should be written to SSD" finally: mgr.close() def test_close_flushes_hot_cache_to_ssd(self, tmp_path): """close() should flush all hot cache entries to SSD.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "wb_flush_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) block_hashes = [] for i in range(3): block_hash = f"wb_flush_blk_{i}__".encode() block_hashes.append(block_hash) cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # No SSD files before close time.sleep(0.3) ssd_files = list((tmp_path / "wb_flush_test").rglob("*.safetensors")) assert len(ssd_files) == 0 # Close flushes to SSD mgr.close() ssd_files = list((tmp_path / "wb_flush_test").rglob("*.safetensors")) assert ( len(ssd_files) == 3 ), f"Expected 3 SSD files after flush, got {len(ssd_files)}" def test_close_flushes_all_blocks_with_small_queue(self, tmp_path): """close() must flush all hot cache blocks even when more blocks exist than the write queue depth. Regression test for #1070: put_nowait() in the shutdown flush loop lost blocks when the bounded write queue filled up faster than the writer thread can drain it. """ queue_depth = 4 block_count = 12 # 3x the queue depth with patch("omlx.cache.paged_ssd_cache._MAX_PENDING_WRITES", queue_depth): mgr = PagedSSDCacheManager( cache_dir=tmp_path / "wb_queue_full_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=10 * 1024**2, ) for i in range(block_count): block_hash = f"wb_qfull_blk_{i:02d}".encode() cache_data = self._make_cache_data() mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test", layer_cache_types=["KVCache"] * 2, ) # All blocks in hot cache, no SSD files yet time.sleep(0.3) ssd_files = list((tmp_path / "wb_queue_full_test").rglob("*.safetensors")) assert len(ssd_files) == 0 mgr.close() ssd_files = list((tmp_path / "wb_queue_full_test").rglob("*.safetensors")) assert len(ssd_files) == block_count, ( f"Expected {block_count} SSD files after flush, got {len(ssd_files)}. " f"Blocks were likely not flushed during shutdown." ) @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestPendingWriteBuffer: """Tests for the pending-write buffer that bridges hot cache eviction to SSD write.""" def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16): return [ ( mx.zeros((1, heads, seq_len, head_dim)), mx.zeros((1, heads, seq_len, head_dim)), ) for _ in range(num_layers) ] def _save_block(self, manager, block_hash, model="test-model"): cache_data = self._make_cache_data() return manager.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name=model, layer_cache_types=["KVCache"] * 2, ) def test_evicted_block_readable_from_pending_buffer(self, tmp_path): """A block evicted from hot cache is still loadable while SSD write is pending.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "pending_buf_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: # Save 2 blocks (fills hot cache) self._save_block(mgr, b"pending_buf_blk0") self._save_block(mgr, b"pending_buf_blk1") # Save 3rd block → evicts block 0 self._save_block(mgr, b"pending_buf_blk2") # Block 0 should still be loadable (from pending buffer, not SSD) with mgr._hot_cache_lock: assert ( b"pending_buf_blk0" not in mgr._hot_cache ), "Block 0 should have been evicted from hot cache" loaded = mgr.load_block(b"pending_buf_blk0") assert ( loaded is not None ), "Evicted block should be readable from pending write buffer" assert len(loaded) == 2 finally: mgr.close() def test_writer_cleanup_empties_buffer(self, tmp_path): """After background writer completes, pending buffer entry is removed.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "writer_cleanup_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: self._save_block(mgr, b"cleanup_test_blk0") self._save_block(mgr, b"cleanup_test_blk1") # Evict block 0 self._save_block(mgr, b"cleanup_test_blk2") # Block 0 should be in pending buffer with mgr._pending_write_hashes_lock: assert b"cleanup_test_blk0" in mgr._pending_write_buffers # Wait for writer to finish time.sleep(1.0) # Buffer should be empty after writer cleanup with mgr._pending_write_hashes_lock: assert b"cleanup_test_blk0" not in mgr._pending_write_buffers assert b"cleanup_test_blk0" not in mgr._pending_write_hashes finally: mgr.close() def test_has_block_sees_pending_entries(self, tmp_path): """has_block() returns True for blocks in the pending write buffer.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "has_block_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: self._save_block(mgr, b"has_block_test_b0") self._save_block(mgr, b"has_block_test_b1") # Evict block 0 self._save_block(mgr, b"has_block_test_b2") with mgr._hot_cache_lock: assert b"has_block_test_b0" not in mgr._hot_cache assert ( mgr.has_block(b"has_block_test_b0") is True ), "has_block should find blocks in the pending write buffer" finally: mgr.close() def test_delete_block_clears_pending_buffer(self, tmp_path): """delete_block() removes the entry from pending write buffer.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "delete_pending_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: self._save_block(mgr, b"del_pending_blk_0") self._save_block(mgr, b"del_pending_blk_1") # Evict block 0 self._save_block(mgr, b"del_pending_blk_2") # Block 0 is in pending buffer with mgr._pending_write_hashes_lock: assert b"del_pending_blk_0" in mgr._pending_write_buffers # Delete it mgr.delete_block(b"del_pending_blk_0") # Should no longer be loadable loaded = mgr.load_block(b"del_pending_blk_0") assert ( loaded is None ), "Deleted block should not be readable from pending buffer" finally: mgr.close() def test_queue_full_inline_fallback_cleans_pending_buffer(self, tmp_path): """Inline fallback removes the block from pending after it reaches SSD.""" import queue as _queue entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "queue_full_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: self._save_block(mgr, b"qf_test_block_00") self._save_block(mgr, b"qf_test_block_01") self._save_block(mgr, b"qf_test_block_02") # This evicts block 1. Force sustained saturation so the bounded # wait expires and inline fallback runs. with patch.object(mgr._write_queue, "put", side_effect=_queue.Full): self._save_block(mgr, b"qf_test_block_03") # Block 1 was written inline and no longer needs pending storage. with mgr._pending_write_hashes_lock: assert ( b"qf_test_block_01" not in mgr._pending_write_buffers ), "Inline-written block should leave the pending buffer" assert ( b"qf_test_block_01" not in mgr._pending_write_hashes ), "Inline-written block should leave pending hashes" stats = mgr.get_stats() assert stats.ssd_write_drops == 0 assert stats.ssd_inline_write_fallbacks == 1 assert mgr.load_block(b"qf_test_block_01") is not None finally: mgr.close() def test_evicted_block_loadable_with_metadata(self, tmp_path): """load_block_with_metadata returns data for pending-buffer blocks.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "meta_load_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: self._save_block(mgr, b"meta_load_blk_00") self._save_block(mgr, b"meta_load_blk_01") # Evict block 0 self._save_block(mgr, b"meta_load_blk_02") cache_data, metadata = mgr.load_block_with_metadata(b"meta_load_blk_00") assert ( cache_data is not None ), "load_block_with_metadata should serve evicted block from pending buffer" assert metadata is not None assert metadata["num_layers"] == 2 assert metadata["token_count"] == 16 assert metadata["model_name"] == "test-model" finally: mgr.close() def test_end_to_end_lifecycle(self, tmp_path): """Full lifecycle: save → evict → pending hit → writer completes → SSD hit.""" entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 mgr = PagedSSDCacheManager( cache_dir=tmp_path / "lifecycle_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: self._save_block(mgr, b"lifecycle_blk_00") self._save_block(mgr, b"lifecycle_blk_01") # Evict block 0 self._save_block(mgr, b"lifecycle_blk_02") # Phase 1: pending buffer hit (before writer completes) loaded = mgr.load_block(b"lifecycle_blk_00") assert loaded is not None, "Should load from pending buffer" assert mgr._stats["hot_cache_hits"] >= 1 # Phase 2: wait for writer to complete time.sleep(1.0) # Buffer should be empty with mgr._pending_write_hashes_lock: assert b"lifecycle_blk_00" not in mgr._pending_write_buffers # Phase 3: SSD hit (block now on disk) loaded_ssd = mgr.load_block(b"lifecycle_blk_00") assert loaded_ssd is not None, "Should load from SSD after writer completes" assert len(loaded_ssd) == 2 finally: mgr.close() @pytest.mark.skipif(not HAS_MLX, reason="MLX not available") class TestSSDWriteBackSaturation: """Queue saturation falls back to inline writes instead of dropping blocks.""" # Mirrors TestPendingWriteBuffer._make_cache_data — small dimensions # so the per-entry footprint is small and predictable for queue tests. def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16): return [ ( mx.zeros((1, heads, seq_len, head_dim)), mx.zeros((1, heads, seq_len, head_dim)), ) for _ in range(num_layers) ] # Mirrors TestPendingWriteBuffer._save_block — uses 2 layers, token_count=16 # so the entry_size formula `2 * 2 * 1 * 2 * 16 * 16 * 4` calibrates Test A. def _save_block(self, manager, block_hash, model="test-model"): cache_data = self._make_cache_data() return manager.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name=model, layer_cache_types=["KVCache"] * 2, ) def test_paged_ssd_cache_stats_default_and_reset(self): """Dataclass: write-back counters default to 0 and reset to 0.""" from omlx.cache.stats import PagedSSDCacheStats # Default is zero. stats = PagedSSDCacheStats() assert stats.ssd_write_drops == 0 assert stats.ssd_inline_write_fallbacks == 0 assert stats.evict_unlink_failures == 0 # reset() returns it to zero from a non-zero state. stats = PagedSSDCacheStats( ssd_write_drops=5, ssd_inline_write_fallbacks=4, evict_unlink_failures=3, saves=2, loads=3, ) stats.reset() assert stats.ssd_write_drops == 0 assert stats.ssd_inline_write_fallbacks == 0 assert stats.evict_unlink_failures == 0 # Verify reset() didn't break the existing fields it already handled. assert stats.saves == 0 assert stats.loads == 0 def test_ssd_write_back_fields_round_trip_through_get_stats(self, tmp_path): """The _stats dict values flow through both stats accessors. Force a non-zero value into _stats so a missing pass-through in either accessor would leave the dataclass at the default 0 and fail this test. The default-0 case is covered by the dataclass-level test in Task 1; this test exercises the wiring specifically. """ mgr = PagedSSDCacheManager( cache_dir=tmp_path / "wiring_test", max_size_bytes=100 * 1024**2, ) try: # Directly poke a non-zero value to test the accessor pass-through # specifically. The real increment sites have their own tests # further down in this class; this test would catch a regression # where a future refactor drops the field from get_stats() or # get_stats_for_model() but leaves the _stats dict key alone. mgr._stats["ssd_write_drops"] = 7 mgr._stats["ssd_inline_write_fallbacks"] = 5 # Global stats accessor. stats = mgr.get_stats() assert stats.ssd_write_drops == 7, ( "get_stats() must pass _stats['ssd_write_drops'] through " "to the dataclass field" ) assert stats.ssd_inline_write_fallbacks == 5, ( "get_stats() must pass _stats['ssd_inline_write_fallbacks'] " "through to the dataclass field" ) # Per-model stats accessor — same wiring, separate code path. # Use any model name; the field is global on the manager. model_stats = mgr.get_stats_for_model("any-model-name") assert model_stats.ssd_write_drops == 7, ( "get_stats_for_model() must pass _stats['ssd_write_drops'] " "through to the dataclass field" ) assert model_stats.ssd_inline_write_fallbacks == 5, ( "get_stats_for_model() must pass " "_stats['ssd_inline_write_fallbacks'] through to the " "dataclass field" ) finally: mgr.close() def test_evict_unlink_failures_round_trips_through_get_stats(self, tmp_path): """evict_unlink_failures is visible through both stats accessors.""" mgr = PagedSSDCacheManager( cache_dir=tmp_path / "unlink_failure_wiring_test", max_size_bytes=100 * 1024**2, ) try: mgr._stats["evict_unlink_failures"] = 3 stats = mgr.get_stats() assert stats.evict_unlink_failures == 3 model_stats = mgr.get_stats_for_model("any-model-name") assert model_stats.evict_unlink_failures == 3 finally: mgr.close() def test_hot_eviction_queue_full_uses_inline_fallback(self, tmp_path): """Site 1: hot-cache eviction → put raises queue.Full → inline write. Patches the real queue's put to raise queue.Full, guaranteeing the fallback path fires on the first eviction without any dependency on the writer thread's drain rate. _enqueue_ssd_write uses put(item, timeout=...) (not put_nowait) so a transient burst can ride over a short writer-backlog window; sustained saturation writes inline. """ import queue as _queue from unittest.mock import patch # entry_size matches the small-dimension helpers above (num_layers=2, # K+V=2 tensors, batch=1, heads=2, seq=16, head_dim=16, float32=4). entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4 max_bytes = entry_size * 2 + 100 # holds exactly 2 entries mgr = PagedSSDCacheManager( cache_dir=tmp_path / "drops_hot_eviction_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=max_bytes, ) try: with patch.object(mgr._write_queue, "put", side_effect=_queue.Full): self._save_block(mgr, b"qf_drop_block_00") self._save_block(mgr, b"qf_drop_block_01") # save_02 evicts block 00 → _enqueue_ssd_write → put raises # queue.Full → inline fallback writes it immediately. self._save_block(mgr, b"qf_drop_block_02") stats = mgr.get_stats() assert stats.ssd_write_drops == 0 assert stats.ssd_inline_write_fallbacks == 1 assert stats.saves_persisted == 1 assert stats.errors == 0 # Block 00 was the one being enqueued when put raised. It should # no longer need the pending buffer because it is already on SSD. with mgr._pending_write_hashes_lock: assert b"qf_drop_block_00" not in mgr._pending_write_buffers assert b"qf_drop_block_00" not in mgr._pending_write_hashes metadata = mgr._index.get(b"qf_drop_block_00") assert metadata is not None assert metadata.file_path.exists() assert mgr.load_block(b"qf_drop_block_00") is not None finally: mgr.close() def test_cold_store_sustained_full_uses_inline_fallback(self, tmp_path): """Site 2: save_block waits on a full queue before writing inline. Hot cache disabled. Even if ``full()`` reports saturation, the cold path must still proceed to ``put(timeout=1.0)`` so transient writer bursts get a chance to drain. Sustained ``queue.Full`` from the put call should write inline and return success. """ import queue as _queue from unittest.mock import patch mgr = PagedSSDCacheManager( cache_dir=tmp_path / "drops_cold_preflight_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=0, # hot cache disabled → cold-store path ) try: calls: list[float | None] = [] def full_put(*args, timeout=None, **kwargs): calls.append(timeout) raise _queue.Full with ( patch.object(mgr._write_queue, "full", return_value=True), patch.object(mgr._write_queue, "put", side_effect=full_put), ): cache_data = self._make_cache_data() block_hash = b"cold_preflight_drop_00" ok = mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test-model", layer_cache_types=["KVCache"] * 2, ) assert ok is True assert calls == [1.0] stats = mgr.get_stats() assert stats.ssd_write_drops == 0 assert stats.ssd_inline_write_fallbacks == 1 assert stats.saves == 1 assert stats.saves_persisted == 1 assert stats.errors == 0 assert mgr._index.contains(block_hash) with mgr._pending_write_hashes_lock: assert block_hash not in mgr._pending_write_hashes assert mgr.load_block(block_hash) is not None finally: mgr.close() def test_cold_store_late_queue_full_uses_inline_fallback(self, tmp_path): """Site 3: save_block put raises queue.Full after the preflight passes. Hot cache disabled. ``put`` is patched to raise queue.Full directly (simulating a sustained writer-backlog saturation that materializes after the preflight check). The block must still be written inline and remain readable. """ import queue as _queue from unittest.mock import patch mgr = PagedSSDCacheManager( cache_dir=tmp_path / "drops_cold_late_exception_test", max_size_bytes=100 * 1024**2, hot_cache_max_bytes=0, ) try: cache_data = self._make_cache_data() block_hash = b"cold_late_drop_00" with patch.object(mgr._write_queue, "put", side_effect=_queue.Full): ok = mgr.save_block( block_hash=block_hash, cache_data=cache_data, token_count=16, model_name="test-model", layer_cache_types=["KVCache"] * 2, ) assert ok is True stats = mgr.get_stats() assert stats.ssd_write_drops == 0 assert stats.ssd_inline_write_fallbacks == 1 assert stats.saves == 1 assert stats.saves_persisted == 1 assert mgr._index.contains(block_hash) with mgr._pending_write_hashes_lock: assert block_hash not in mgr._pending_write_hashes assert mgr.load_block(block_hash) is not None finally: mgr.close()