amberjs 1.17.0

JavaScript/TypeScript runtime built with Rust and V8
Documentation
// 模型缓存系统
// 提供高效的 AI 模型缓存和管理

use anyhow::Result;
use std::collections::{BTreeMap, HashMap};
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};

/// 缓存策略
#[derive(Debug, Clone)]
pub enum CacheStrategy {
    /// 最近最少使用
    LRU,
    /// 最不经常使用
    LFU,
    /// 先进先出
    FIFO,
    /// 智能缓存(基于访问模式)
    Smart,
}
/// 缓存统计
#[derive(Debug, Clone)]
pub struct CacheStats {
    pub hits: u64,
    pub misses: u64,
    pub hit_rate: f64,
    pub total_models: usize,
    pub total_size: usize,
    pub max_size: usize,
}
/// 缓存条目
#[derive(Debug)]
struct CacheEntry {
    model: AIModel,
    access_count: u64,
    last_access: Instant,
    created_at: Instant,
    size: usize,
}
/// 模型缓存
#[derive(Debug)]
pub struct ModelCache {
    models: Arc<RwLock<HashMap<String, CacheEntry>>>,
    max_size: usize,
    strategy: CacheStrategy,
    stats: Arc<RwLock<CacheStats>>,
}
impl ModelCache {
    /// 创建新的模型缓存
    pub async fn new(max_size: usize) -> Result<Self> {
        let cache: _ = ModelCache {
            models: Arc::new(Mutex::new(HashMap::new())),
            max_size,
            strategy: CacheStrategy::Smart,
            stats: Arc::new(Mutex::new(CacheStats {
                hits: 0,
                misses: 0,
                hit_rate: 0.0,
                total_models: 0,
                total_size: 0,
                max_size,
            })),
        };
        // 启动缓存清理任务
        cache.start_cleanup_task().await?;
        Ok(cache)
    }
    /// 获取模型
    pub async fn get(&self, model_id: &str) -> Result<Option<AIModel>> {
        let mut models = self.models.write().await;
        let mut hit = false;
        let model_count: _ = models.len();
        let total_size: _ = models.values().map(|e| e.size).sum();
        if let Some(entry) = models.get_mut(model_id) {
            // 更新访问统计
            entry.access_count += 1;
            entry.last_access = Instant::now();
            hit = true;
            let model: _ = entry.model.clone();
            drop(models); // 释放锁
                          // 更新缓存统计
            {
                let mut stats = self.stats.write().await;
                stats.hits += 1;
                stats.total_models = model_count;
                stats.total_size = total_size;
                stats.hit_rate = stats.hits as f64 / (stats.hits + stats.misses) as f64;
            }
            return Ok(Some(model));
        }
        drop(models); // 释放锁
                      // 缓存未命中
        {
            let mut stats = self.stats.write().await;
            stats.misses += 1;
            stats.total_models = model_count;
            stats.total_size = total_size;
            stats.hit_rate = stats.hits as f64 / (stats.hits + stats.misses) as f64;
        }
        Ok(None)
    }
    /// 放入模型
    pub async fn put(&self, model_id: String, model: AIModel) -> Result<()> {
        let mut models = self.models.write().await;
        // 计算模型大小(简化实现)
        let size: _ = self.calculate_model_size(&model);
        // 如果缓存已满,触发清理
        if models.len() >= self.max_size {
            self.evict_if_needed(&mut models).await?;
        }
        // 插入或更新模型
        let entry: _ = CacheEntry {
            model,
            access_count: 1,
            last_access: Instant::now(),
            created_at: Instant::now(),
            size,
        };
        models.insert(model_id, entry);
        // 更新统计
        {
            let mut stats = self.stats.write().await;
            stats.total_models = models.len();
            stats.total_size = models.values().map(|e| e.size).sum();
        }
        Ok(())
    }
    /// 预加载模型
    pub async fn preload(&self, model_ids: &[String]) -> Result<()> {
        let loader: _ = ModelLoader::new();
        for model_id in model_ids {
            // 检查是否已缓存
            if self.get(model_id).await?.is_some() {
                continue;
            }
            // 加载模型
            let model: _ = loader.load(model_id).await?;
            self.put(model_id.clone(), model).await?;
        }
        Ok(())
    }
    /// 智能预取
    pub async fn smart_prefetch(&self, access_pattern: &[String]) -> Result<()> {
        // 分析访问模式,预测未来需要的模型
        let predictions: _ = self.predict_future_accesses(access_pattern).await?;
        // 预加载预测的模型
        if !predictions.is_empty() {
            self.preload(&predictions).await?;
        }
        Ok(())
    }
    /// 预测未来访问
    async fn predict_future_accesses(&self, access_pattern: &[String]) -> Result<Vec<String>> {
        let mut predictions = Vec::new();
        // 简化的预测算法
        // 实际实现中会使用机器学习模型来预测
        if access_pattern.len() >= 2 {
            // 如果访问了 A 和 B,预测可能访问 C
            let last_two: _ = &access_pattern[access_pattern.len() - 2..];
            if last_two[0] == "bert" && last_two[1] == "gpt" {
                predictions.push("resnet50".to_string());
            } else if last_two[0] == "gpt" && last_two[1] == "resnet50" {
                predictions.push("bert".to_string());
            }
        }
        Ok(predictions)
    }
    /// 清理过期模型
    pub async fn cleanup_expired(&self, max_age: Duration) -> Result<usize> {
        let mut models = self.models.write().await;
        let mut removed_count = 0;
        let now: _ = Instant::now();
        let to_remove: Vec<String> = models
            .values()
            .filter(|entry| now.duration_since(entry.created_at) > max_age)
            .filter(|entry| entry.access_count == 1) // 只移除未访问的
            .map(|entry| entry.model.id.clone())
            .collect();
        for model_id in to_remove {
            models.remove(&model_id);
            removed_count += 1;
        }
        // 更新统计
        {
            let mut stats = self.stats.write().await;
            stats.total_models = models.len();
            stats.total_size = models.values().map(|e| e.size).sum();
        }
        Ok(removed_count)
    }
    /// 启动清理任务
    async fn start_cleanup_task(&self) -> Result<()> {
        let models_ref: _ = self.models.clone();
        let stats_ref: _ = self.stats.clone();
        let max_age: _ = Duration::from_secs(300); // 5分钟
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(60)); // 每分钟清理一次
            loop {
                interval.tick().await;
                let mut models = models_ref.write().await;
                let now: _ = Instant::now();
                // 清理过期模型
                let to_remove: Vec<String> = models
                    .values()
                    .filter(|entry| now.duration_since(entry.created_at) > max_age)
                    .filter(|entry| entry.access_count == 1)
                    .map(|entry| entry.model.id.clone())
                    .collect();
                for model_id in to_remove {
                    models.remove(&model_id);
                }
                // 更新统计
                {
                    let mut stats = stats_ref.write().await;
                    stats.total_models = models.len();
                    stats.total_size = models.values().map(|e| e.size).sum();
                }
            }
        });
        Ok(())
    }
    /// 需要时驱逐模型
    async fn evict_if_needed(&self, models: &mut HashMap<String, CacheEntry>) -> Result<()> {
        if models.len() < self.max_size {
            return Ok(());
        }
        // 根据策略驱逐模型
        match self.strategy {
            CacheStrategy::LRU => self.evict_lru(models).await?,
            CacheStrategy::LFU => self.evict_lfu(models).await?,
            CacheStrategy::FIFO => self.evict_fifo(models).await?,
            CacheStrategy::Smart => self.evict_smart(models).await?,
        }
        Ok(())
    }
    /// LRU 驱逐
    async fn evict_lru(&self, models: &mut HashMap<String, CacheEntry>) -> Result<()> {
        let lru_key: _ = models
            .iter()
            .min_by_key(|(_, entry)| entry.last_access)
            .map(|(key, _)| key.clone());
        if let Some(key) = lru_key {
            models.remove(&key);
        }
        Ok(())
    }
    /// LFU 驱逐
    async fn evict_lfu(&self, models: &mut HashMap<String, CacheEntry>) -> Result<()> {
        let lfu_key: _ = models
            .iter()
            .min_by_key(|(_, entry)| entry.access_count)
            .map(|(key, _)| key.clone());
        if let Some(key) = lfu_key {
            models.remove(&key);
        }
        Ok(())
    }
    /// FIFO 驱逐
    async fn evict_fifo(&self, models: &mut HashMap<String, CacheEntry>) -> Result<()> {
        let fifo_key: _ = models
            .iter()
            .min_by_key(|(_, entry)| entry.created_at)
            .map(|(key, _)| key.clone());
        if let Some(key) = fifo_key {
            models.remove(&key);
        }
        Ok(())
    }
    /// 智能驱逐
    async fn evict_smart(&self, models: &mut HashMap<String, CacheEntry>) -> Result<()> {
        // 综合考虑访问频率、时间和大小
        let smart_key: _ = models
            .iter()
            .min_by(|a, b| {
                let score_a: _ = self.calculate_eviction_score(a.1);
                let score_b: _ = self.calculate_eviction_score(b.1);
                score_a
                    .partial_cmp(&score_b)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .map(|(key, _)| key.clone());
        if let Some(key) = smart_key {
            models.remove(&key);
        }
        Ok(())
    }
    /// 计算驱逐分数
    fn calculate_eviction_score(&self, entry: &CacheEntry) -> f64 {
        let recency: _ = entry.last_access.elapsed().as_secs_f64();
        let frequency: _ = entry.access_count as f64;
        let age: _ = entry.created_at.elapsed().as_secs_f64();
        // 分数越低越容易被驱逐
        (recency + age) / (frequency + 1.0)
    }
    /// 计算模型大小
    fn calculate_model_size(&self, model: &AIModel) -> usize {
        // 简化实现:基于输入输出维度计算大小
        let input_size: usize = model.input_shape.iter().product();
        let output_size: usize = model.output_shape.iter().product();
        input_size + output_size + model.parameters.len()
    }
    /// 获取缓存统计
    pub async fn get_stats(&self) -> Result<CacheStats> {
        let stats: _ = self.stats.read().await;
        Ok(stats.clone())
    }
    /// 清空缓存
    pub async fn clear(&self) -> Result<()> {
        let mut models = self.models.write().await;
        models.clear();
        // 重置统计
        {
            let mut stats = self.stats.write().await;
            stats.total_models = 0;
            stats.total_size = 0;
        }
        Ok(())
    }
    /// 获取缓存大小
    pub async fn size(&self) -> Result<usize> {
        let models: _ = self.models.read().await;
        Ok(models.len())
    }
    /// 检查是否包含模型
    pub async fn contains(&self, model_id: &str) -> Result<bool> {
        let models: _ = self.models.read().await;
        Ok(models.contains_key(model_id))
    }
}