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 {
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); 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(())
}
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(())
}
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(())
}
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))
}
}