rustytorch/crates/production/rtx-inference/tests/cache_test.rs

//! Comprehensive tests for paged KV cache functionality

use rtx_inference::{
    InferenceError,
    cache::{EvictionPolicy, KvCacheConfig, MemoryTier, PagedKvCache},
};
use rtx_tensor::{Device, Tensor};
use std::time::Duration;

#[cfg(test)]
mod cache_tests {
    use super::*;

    fn create_test_config() -> KvCacheConfig {
        KvCacheConfig {
            page_size: 4,
            max_pages_gpu: 10,
            max_pages_cpu: 20,
            max_pages_nvme: 50,
            eviction_policy: EvictionPolicy::LRU,
            prefetch_enabled: false,
            prefetch_distance: 4,
            compression_enabled: false,
            compression_ratio_threshold: 2.0,
            memory_pressure_threshold: 0.8,
            model_isolation: true,
            persistence_enabled: false,
            persistence_path: "/tmp/rtx_cache".to_string(),
            ..Default::default()
        }
    }

    fn create_test_tensor(device: &Device) -> Tensor {
        Tensor::from_data(vec![1.0, 2.0, 3.0, 4.0], vec![4, 1], device).unwrap()
    }

    #[tokio::test]
    async fn test_cache_creation() {
        let config = create_test_config();
        let device = Device::cpu();
        let cache = PagedKvCache::new(config, device).await;
        assert!(cache.is_ok());
    }

    #[tokio::test]
    async fn test_page_allocation() {
        let config = create_test_config();
        let device = Device::cpu();
        let mut cache = PagedKvCache::new(config, device).await.unwrap();

        let pages = cache.allocate_pages("test_seq", 10).await.unwrap();
        assert_eq!(pages.len(), 3); // 10 tokens / 4 tokens per page = 3 pages
    }

    #[tokio::test]
    async fn test_store_and_retrieve_kv_data() {
        let config = create_test_config();
        let device = Device::cpu();
        let mut cache = PagedKvCache::new(config, device.clone()).await.unwrap();

        // Allocate pages
        let _pages = cache.allocate_pages("test_seq", 8).await.unwrap();

        // Create KV tensors
        let key_data = create_test_tensor(&device);
        let value_data = create_test_tensor(&device);

        // Store data
        let result = cache
            .store_kv_data("test_seq", 0, &key_data, &value_data)
            .await;
        assert!(result.is_ok());

        // Retrieve data
        let (retrieved_keys, retrieved_values) = cache.get_kv_data("test_seq", 0, 4).await.unwrap();
        assert_eq!(retrieved_keys.shape(), key_data.shape());
        assert_eq!(retrieved_values.shape(), value_data.shape());
    }

    #[tokio::test]
    #[ignore = "Pre-existing tier migration assertion failure"]
    async fn test_tier_migration() {
        let config = create_test_config();
        let device = Device::cpu();
        let mut cache = PagedKvCache::new(config, device.clone()).await.unwrap();

        // Allocate in GPU tier
        let pages = cache
            .allocate_pages_in_tier("test_seq", 4, MemoryTier::GPU)
            .await
            .unwrap();
        assert_eq!(pages.len(), 1);

        let page_info = cache.get_page_info(pages[0]).await.unwrap();
        assert_eq!(page_info.tier, MemoryTier::GPU);

        // Migrate to CPU
        cache
            .migrate_to_tier("test_seq", MemoryTier::CPU)
            .await
            .unwrap();

        let page_info_after = cache.get_page_info(pages[0]).await.unwrap();
        assert_eq!(page_info_after.tier, MemoryTier::CPU);
    }

    #[tokio::test]
    async fn test_cache_stats() {
        let config = create_test_config();
        let device = Device::cpu();
        let mut cache = PagedKvCache::new(config, device).await.unwrap();

        // Allocate some pages
        cache.allocate_pages("seq1", 8).await.unwrap();
        cache.allocate_pages("seq2", 4).await.unwrap();

        let stats = cache.get_stats().await;
        assert!(stats.total_pages_allocated >= 3);
        assert!(stats.gpu_pages_used >= 3);
    }

    #[tokio::test]
    async fn test_eviction_policies() {
        for policy in [
            EvictionPolicy::LRU,
            EvictionPolicy::LFU,
            EvictionPolicy::FIFO,
            EvictionPolicy::Random,
        ] {
            let config = KvCacheConfig {
                eviction_policy: policy,
                ..create_test_config()
            };
            let device = Device::cpu();
            let cache = PagedKvCache::new(config, device).await;
            assert!(
                cache.is_ok(),
                "Failed to create cache with policy {:?}",
                policy
            );
        }
    }

    #[tokio::test]
    async fn test_compression_config() {
        let config = KvCacheConfig {
            compression_enabled: true,
            compression_ratio_threshold: 2.0,
            ..create_test_config()
        };
        let device = Device::cpu();
        let mut cache = PagedKvCache::new(config, device.clone()).await.unwrap();

        // Allocate in CPU tier (where compression applies)
        cache
            .allocate_pages_in_tier("test_seq", 4, MemoryTier::CPU)
            .await
            .unwrap();

        let key_data = create_test_tensor(&device);
        let value_data = create_test_tensor(&device);

        cache
            .store_kv_data("test_seq", 0, &key_data, &value_data)
            .await
            .unwrap();

        // Compression stats should be updated
        let stats = cache.get_stats().await;
        assert!(stats.compression_ratio > 0.0);
    }

    #[tokio::test]
    async fn test_persistence_config() {
        let config = KvCacheConfig {
            persistence_enabled: true,
            persistence_path: "/tmp/test_cache".to_string(),
            ..create_test_config()
        };
        let device = Device::cpu();
        let cache = PagedKvCache::new(config, device).await.unwrap();

        let result = cache.persist_to_disk().await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_cache_recovery() {
        let config = KvCacheConfig {
            persistence_enabled: true,
            persistence_path: "/tmp/test_cache_recovery".to_string(),
            ..create_test_config()
        };
        let device = Device::cpu();
        let mut cache = PagedKvCache::new(config, device).await.unwrap();

        let result = cache.recover_from_disk().await;
        assert!(result.is_ok());
    }
}