use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct ModelCacheConfig {
pub max_memory_mb: usize,
pub max_disk_gb: usize,
pub enable_compression: bool,
pub enable_prefetch: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub enum CacheTier {
L1Memory,
L2Disk,
L3Remote,
}
#[derive(Debug, Clone)]
struct CacheEntry {
model_name: String,
model_data: Arc<ModelData>,
tier: CacheTier,
last_access: Instant,
access_count: u64,
size_bytes: usize,
}
#[derive(Debug, Clone)]
struct ModelData {
weights: Vec<u8>,
metadata: ModelMetadata,
}
#[derive(Debug, Clone)]
struct ModelMetadata {
model_name: String,
version: String,
size_mb: usize,
compression_ratio: f32,
load_time: Duration,
}
#[derive(Debug, Clone)]
struct AccessPattern {
model_name: String,
access_count: u64,
last_access: Instant,
access_intervals: Vec<Duration>,
}
pub struct ModelCache {
config: ModelCacheConfig,
l1_cache: Arc<RwLock<HashMap<String, CacheEntry>>>,
l2_cache: Arc<Mutex<HashMap<String, CacheEntry>>>,
l3_cache: Arc<Mutex<HashMap<String, CacheEntry>>>,
access_patterns: Arc<RwLock<HashMap<String, AccessPattern>>>,
stats: Arc<Mutex<CacheStats>>,
}
#[derive(Debug, Clone, Default)]
pub struct CacheStats {
pub l1_hits: u64,
pub l2_hits: u64,
pub l3_hits: u64,
pub misses: u64,
pub total_requests: u64,
pub memory_usage_mb: usize,
pub disk_usage_mb: usize,
}
#[derive(Debug, Clone)]
pub struct CacheResult {
pub model_data: Arc<ModelData>,
pub tier: CacheTier,
pub load_time: Duration,
pub from_cache: bool,
}
impl ModelCache {
pub fn new(config: ModelCacheConfig) -> Result<Self, String> {
Ok(ModelCache {
config: config.clone(),
l1_cache: Arc::new(Mutex::new(HashMap::new()))
l2_cache: Arc::new(Mutex::new(HashMap::new()))
l3_cache: Arc::new(Mutex::new(HashMap::new()))
access_patterns: Arc::new(Mutex::new(HashMap::new()))
stats: Arc::new(Mutex::new(CacheStats::default()))
})
}
pub fn load_model(&mut self, model_name: String) -> Result<CacheResult, String> {
let start_time: _ = Instant::now();
if let Some(entry) = self.get_from_l1(&model_name) {
self.update_stats(entry.tier.clone(), true);
return Ok(CacheResult {
model_data: entry.model_data.clone(),
tier: entry.tier.clone(),
load_time: Duration::from_millis(1),
from_cache: true,
});
}
if let Some(entry) = self.get_from_l2(&model_name) {
self.promote_to_l1(entry.clone())?;
self.update_stats(CacheTier::L2Disk, true);
return Ok(CacheResult {
model_data: entry.model_data.clone(),
tier: CacheTier::L2Disk,
load_time: Duration::from_millis(5),
from_cache: true,
});
}
if let Some(entry) = self.get_from_l3(&model_name) {
self.promote_to_l2(entry.clone())?;
self.update_stats(CacheTier::L3Remote, true);
return Ok(CacheResult {
model_data: entry.model_data.clone(),
tier: CacheTier::L3Remote,
load_time: Duration::from_millis(20),
from_cache: true,
});
}
let model_data: _ = self.load_from_disk(&model_name)?;
let load_time: _ = start_time.elapsed();
let entry: _ = CacheEntry {
model_name: model_name.clone(),
model_data: Arc::new(Mutex::new(model_data)))
tier: CacheTier::L2Disk,
last_access: Instant::now(),
access_count: 1,
size_bytes: 1024 * 1024, };
self.store_to_l2(entry.clone())?;
self.update_stats(CacheTier::L2Disk, false);
Ok(CacheResult {
model_data: entry.model_data.clone(),
tier: CacheTier::L2Disk,
load_time,
from_cache: false,
})
}
fn get_from_l1(&self, model_name: &str) -> Option<CacheEntry> {
let cache: _ = self.l1_cache.read().unwrap();
if let Some(entry) = cache.get(model_name).cloned() {
drop(cache);
self.record_access(model_name);
Some(entry)
} else {
None
}
}
fn get_from_l2(&self, model_name: &str) -> Option<CacheEntry> {
let cache: _ = self.l2_cache.lock().unwrap();
if let Some(entry) = cache.get(model_name).cloned() {
drop(cache);
self.record_access(model_name);
Some(entry)
} else {
None
}
}
fn get_from_l3(&self, model_name: &str) -> Option<CacheEntry> {
let cache: _ = self.l3_cache.lock().unwrap();
if let Some(entry) = cache.get(model_name).cloned() {
drop(cache);
self.record_access(model_name);
Some(entry)
} else {
None
}
}
fn promote_to_l1(&self, entry: CacheEntry) -> Result<(), String> {
let mut l1_cache = self.l1_cache.write().unwrap();
if self.get_memory_usage() > self.config.max_memory_mb {
self.evict_from_l1()?;
}
l1_cache.insert(entry.model_name.clone(), entry);
Ok(())
}
fn promote_to_l2(&self, entry: CacheEntry) -> Result<(), String> {
let mut l2_cache = self.l2_cache.lock().unwrap();
if self.get_disk_usage() > self.config.max_disk_gb * 1024 {
self.evict_from_l2()?;
}
l2_cache.insert(entry.model_name.clone(), entry);
Ok(())
}
fn store_to_l2(&self, entry: CacheEntry) -> Result<(), String> {
let mut l2_cache = self.l2_cache.lock().unwrap();
l2_cache.insert(entry.model_name.clone(), entry);
Ok(())
}
fn load_from_disk(&self, model_name: &str) -> Result<ModelData, String> {
let weights: _ = vec![0u8; 1024 * 1024]; let metadata: _ = ModelMetadata {
model_name: model_name.to_string(),
version: "1.0".to_string(),
size_mb: 1,
compression_ratio: 0.8,
load_time: Duration::from_millis(100),
};
Ok(ModelData { weights, metadata })
}
fn record_access(&self, model_name: &str) {
let mut patterns = self.access_patterns.write().unwrap();
let now: _ = Instant::now();
if let Some(pattern) = patterns.get_mut(model_name) {
pattern.access_count += 1;
let interval: _ = now.duration_since(pattern.last_access);
pattern.access_intervals.push(interval);
pattern.last_access = now;
if pattern.access_intervals.len() > 100 {
pattern.access_intervals.remove(0);
}
} else {
patterns.insert(
model_name.to_string(),
AccessPattern {
model_name: model_name.to_string(),
access_count: 1,
last_access: now,
access_intervals: Vec::new(),
},
);
}
}
pub fn preload_model(&mut self, model_name: String) -> Result<(), String> {
let patterns: _ = self.access_patterns.read().unwrap();
if let Some(pattern) = patterns.get(&model_name) {
if pattern.access_count > 10 {
drop(patterns);
let _: _ = self.load_model(model_name);
}
}
Ok(())
}
pub fn get_prefetch_recommendations(&self) -> Vec<String> {
let patterns: _ = self.access_patterns.read().unwrap();
let mut recommendations = Vec::new();
for (model_name, pattern) in patterns.iter() {
if pattern.access_count > 20 {
let avg_interval: _ = pattern
.access_intervals
.iter()
.sum::<Duration>()
.div_f32(pattern.access_intervals.len() as f32);
let now: _ = Instant::now();
let time_since_last: _ = now.duration_since(pattern.last_access);
if time_since_last < avg_interval {
recommendations.push(model_name.clone());
}
}
}
recommendations.sort_by(|a, b| {
let count_a: _ = patterns.get(a).map(|p| p.access_count).unwrap_or(0);
let count_b: _ = patterns.get(b).map(|p| p.access_count).unwrap_or(0);
count_b.cmp(&count_a)
});
recommendations.into_iter().take(10).collect()
}
fn evict_from_l1(&self) -> Result<(), String> {
let mut cache = self.l1_cache.write().unwrap();
if cache.is_empty() {
return Ok(());
}
let oldest_key: _ = cache
.iter()
.min_by_key(|(_, entry)| entry.last_access)
.map(|(key, _)| key.clone())
.unwrap();
cache.remove(&oldest_key);
Ok(())
}
fn evict_from_l2(&self) -> Result<(), String> {
let mut cache = self.l2_cache.lock().unwrap();
if cache.is_empty() {
return Ok(());
}
let least_frequent_key: _ = cache
.iter()
.min_by_key(|(_, entry)| entry.access_count)
.map(|(key, _)| key.clone())
.unwrap();
cache.remove(&least_frequent_key);
Ok(())
}
fn update_stats(&self, tier: CacheTier, hit: bool) {
let mut stats = self.stats.lock().unwrap();
stats.total_requests += 1;
match (tier, hit) {
(CacheTier::L1Memory, true) => stats.l1_hits += 1,
(CacheTier::L2Disk, true) => stats.l2_hits += 1,
(CacheTier::L3Remote, true) => stats.l3_hits += 1,
(_, false) => stats.misses += 1,
_ => {}
}
}
pub fn get_hit_rate(&self) -> f64 {
let stats: _ = self.stats.lock().unwrap();
if stats.total_requests == 0 {
return 0.0;
}
let hits: _ = stats.l1_hits + stats.l2_hits + stats.l3_hits;
hits as f64 / stats.total_requests as f64
}
pub fn get_prefetch_accuracy(&self) -> f64 {
let patterns: _ = self.access_patterns.read().unwrap();
let mut accurate_predictions = 0;
let mut total_predictions = 0;
for (_model_name, pattern) in patterns.iter() {
total_predictions += 1;
if pattern.access_count > 10 {
accurate_predictions += 1;
}
}
if total_predictions == 0 {
0.0
} else {
accurate_predictions as f64 / total_predictions as f64
}
}
pub fn get_memory_usage(&self) -> usize {
let cache: _ = self.l1_cache.read().unwrap();
let mut usage = 0;
for entry in cache.values() {
usage += entry.size_bytes;
}
usage / (1024 * 1024) }
pub fn get_disk_usage(&self) -> usize {
let cache: _ = self.l2_cache.lock().unwrap();
let mut usage = 0;
for entry in cache.values() {
usage += entry.size_bytes;
}
usage / (1024 * 1024) }
pub fn optimize(&self) {
let now: _ = Instant::now();
let mut l1_cache = self.l1_cache.write().unwrap();
l1_cache.retain(|_, entry| {
now.duration_since(entry.last_access) < Duration::from_secs(3600)
});
let mut l2_cache = self.l2_cache.lock().unwrap();
l2_cache.retain(|_, entry| {
now.duration_since(entry.last_access) < Duration::from_secs(7200)
});
}
pub fn force_gc(&self) {
let mut l1_cache = self.l1_cache.write().unwrap();
l1_cache.clear();
let mut l2_cache = self.l2_cache.lock().unwrap();
l2_cache.clear();
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_cache_creation() {
let config: _ = ModelCacheConfig {
max_memory_mb: 1024,
max_disk_gb: 10,
enable_compression: true,
enable_prefetch: true,
};
let cache: _ = ModelCache::new(config);
assert!(cache.is_ok());
}
#[test]
fn test_model_loading() {
let mut cache = ModelCache::new(ModelCacheConfig {
max_memory_mb: 1024,
max_disk_gb: 10,
enable_compression: true,
enable_prefetch: false,
}).unwrap();
let result: _ = cache.load_model("test-model".to_string());
assert!(result.is_ok());
}
#[test]
fn test_cache_hit_rate() {
let mut cache = ModelCache::new(ModelCacheConfig {
max_memory_mb: 1024,
max_disk_gb: 10,
enable_compression: true,
enable_prefetch: true,
}).unwrap();
let _: _ = cache.load_model("model-1".to_string());
for _ in 0..10 {
let _: _ = cache.load_model("model-1".to_string());
}
let hit_rate: _ = cache.get_hit_rate();
println!("Hit rate: {}%", hit_rate * 100.0);
assert!(hit_rate > 0.8);
}
#[test]
fn test_prefetch_recommendations() {
let mut cache = ModelCache::new(ModelCacheConfig {
max_memory_mb: 1024,
max_disk_gb: 10,
enable_compression: true,
enable_prefetch: true,
}).unwrap();
for _ in 0..30 {
let _: _ = cache.load_model("popular-model".to_string());
}
let recommendations: _ = cache.get_prefetch_recommendations();
assert!(!recommendations.is_empty());
}
}