Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,339 @@
|
||||
//! Eviction policies and cache persistence.
|
||||
|
||||
use super::paged_kv_cache::PagedKvCache;
|
||||
use super::types::{EvictionPolicy, MemoryTier, PageId};
|
||||
use crate::error::{InferenceError, InferenceResult};
|
||||
use std::time::Instant;
|
||||
use tracing::{info, trace, warn};
|
||||
|
||||
impl PagedKvCache {
|
||||
/// Evict a page using configured policy
|
||||
pub async fn evict_page(&self) -> InferenceResult<()> {
|
||||
match self.config.eviction_policy {
|
||||
EvictionPolicy::LRU => self.evict_lru_page().await,
|
||||
EvictionPolicy::LFU => self.evict_lfu_page().await,
|
||||
EvictionPolicy::FIFO => self.evict_fifo_page().await,
|
||||
EvictionPolicy::Random => self.evict_random_page().await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Evict least recently used page
|
||||
pub(crate) async fn evict_lru_page(&self) -> InferenceResult<()> {
|
||||
let lru_page_id = {
|
||||
let lru_tracker = self.lru_tracker.lock().await;
|
||||
lru_tracker.get_lru()
|
||||
};
|
||||
|
||||
if let Some(page_id) = lru_page_id {
|
||||
self.remove_page(page_id).await?;
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.evictions_lru += 1;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Evict least frequently used page
|
||||
pub(crate) async fn evict_lfu_page(&self) -> InferenceResult<()> {
|
||||
let mut lfu_page_id = None;
|
||||
let mut min_access_count = usize::MAX;
|
||||
|
||||
{
|
||||
let gpu_pages = self.gpu_pages.read().await;
|
||||
for (page_id, page) in gpu_pages.iter() {
|
||||
if page.access_count < min_access_count {
|
||||
min_access_count = page.access_count;
|
||||
lfu_page_id = Some(*page_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if lfu_page_id.is_none() {
|
||||
let cpu_pages = self.cpu_pages.read().await;
|
||||
for (page_id, page) in cpu_pages.iter() {
|
||||
if page.access_count < min_access_count {
|
||||
min_access_count = page.access_count;
|
||||
lfu_page_id = Some(*page_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if lfu_page_id.is_none() {
|
||||
let nvme_pages = self.nvme_pages.read().await;
|
||||
for (page_id, page) in nvme_pages.iter() {
|
||||
if page.access_count < min_access_count {
|
||||
lfu_page_id = Some(*page_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(page_id) = lfu_page_id {
|
||||
self.remove_page(page_id).await?;
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.evictions_lfu += 1;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Evict first in, first out page
|
||||
pub(crate) async fn evict_fifo_page(&self) -> InferenceResult<()> {
|
||||
let mut oldest_page_id = None;
|
||||
let mut oldest_time = Instant::now();
|
||||
let mut evict_from_tier = MemoryTier::GPU;
|
||||
|
||||
{
|
||||
let gpu_pages = self.gpu_pages.read().await;
|
||||
for (page_id, page) in gpu_pages.iter() {
|
||||
if page.created_at < oldest_time {
|
||||
oldest_time = page.created_at;
|
||||
oldest_page_id = Some(*page_id);
|
||||
evict_from_tier = MemoryTier::GPU;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let cpu_pages = self.cpu_pages.read().await;
|
||||
for (page_id, page) in cpu_pages.iter() {
|
||||
if page.created_at < oldest_time {
|
||||
oldest_time = page.created_at;
|
||||
oldest_page_id = Some(*page_id);
|
||||
evict_from_tier = MemoryTier::CPU;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let nvme_pages = self.nvme_pages.read().await;
|
||||
for (page_id, page) in nvme_pages.iter() {
|
||||
if page.created_at < oldest_time {
|
||||
oldest_page_id = Some(*page_id);
|
||||
evict_from_tier = MemoryTier::NVME;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(page_id) = oldest_page_id {
|
||||
match evict_from_tier {
|
||||
MemoryTier::GPU => {
|
||||
self.gpu_pages.write().await.remove(&page_id);
|
||||
}
|
||||
MemoryTier::CPU => {
|
||||
self.cpu_pages.write().await.remove(&page_id);
|
||||
}
|
||||
MemoryTier::NVME => {
|
||||
self.nvme_pages.write().await.remove(&page_id);
|
||||
}
|
||||
}
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.evictions_due_to_pressure += 1;
|
||||
trace!(
|
||||
"FIFO evicted page {} from {:?} tier",
|
||||
page_id, evict_from_tier
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Evict random page
|
||||
pub(crate) async fn evict_random_page(&self) -> InferenceResult<()> {
|
||||
let mut rng = fastrand::Rng::new();
|
||||
let mut all_pages = Vec::new();
|
||||
|
||||
{
|
||||
let gpu_pages = self.gpu_pages.read().await;
|
||||
for page_id in gpu_pages.keys() {
|
||||
all_pages.push((*page_id, MemoryTier::GPU));
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let cpu_pages = self.cpu_pages.read().await;
|
||||
for page_id in cpu_pages.keys() {
|
||||
all_pages.push((*page_id, MemoryTier::CPU));
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let nvme_pages = self.nvme_pages.read().await;
|
||||
for page_id in nvme_pages.keys() {
|
||||
all_pages.push((*page_id, MemoryTier::NVME));
|
||||
}
|
||||
}
|
||||
|
||||
if !all_pages.is_empty() {
|
||||
let random_idx = rng.usize(0..all_pages.len());
|
||||
let (page_id, tier) = all_pages[random_idx];
|
||||
|
||||
match tier {
|
||||
MemoryTier::GPU => {
|
||||
self.gpu_pages.write().await.remove(&page_id);
|
||||
}
|
||||
MemoryTier::CPU => {
|
||||
self.cpu_pages.write().await.remove(&page_id);
|
||||
}
|
||||
MemoryTier::NVME => {
|
||||
self.nvme_pages.write().await.remove(&page_id);
|
||||
}
|
||||
}
|
||||
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.evictions_due_to_pressure += 1;
|
||||
trace!("Random evicted page {} from {:?} tier", page_id, tier);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove page from cache
|
||||
pub(crate) async fn remove_page(&self, page_id: PageId) -> InferenceResult<()> {
|
||||
let tier = {
|
||||
let page_info = self.page_info.read().await;
|
||||
page_info.get(&page_id).map(|info| info.tier)
|
||||
};
|
||||
|
||||
if let Some(tier) = tier {
|
||||
match tier {
|
||||
MemoryTier::GPU => {
|
||||
self.gpu_pages.write().await.remove(&page_id);
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.gpu_pages_used = stats.gpu_pages_used.saturating_sub(1);
|
||||
}
|
||||
MemoryTier::CPU => {
|
||||
self.cpu_pages.write().await.remove(&page_id);
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.cpu_pages_used = stats.cpu_pages_used.saturating_sub(1);
|
||||
}
|
||||
MemoryTier::NVME => {
|
||||
self.nvme_pages.write().await.remove(&page_id);
|
||||
let mut stats = self.stats.write().await;
|
||||
stats.nvme_pages_used = stats.nvme_pages_used.saturating_sub(1);
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
self.page_info.write().await.remove(&page_id);
|
||||
}
|
||||
|
||||
{
|
||||
let mut lru_tracker = self.lru_tracker.lock().await;
|
||||
lru_tracker.remove(page_id);
|
||||
}
|
||||
|
||||
trace!("Removed page {} from tier {}", page_id, tier);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Persist cache to disk
|
||||
pub async fn persist_to_disk(&self) -> InferenceResult<()> {
|
||||
if !self.config.persistence_enabled {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
use std::fs;
|
||||
|
||||
let cache_dir = std::path::Path::new("./cache_persistence");
|
||||
if !cache_dir.exists() {
|
||||
fs::create_dir_all(cache_dir).map_err(|e| {
|
||||
InferenceError::kv_cache_error(
|
||||
"persist",
|
||||
format!("Failed to create cache directory: {e}"),
|
||||
)
|
||||
})?;
|
||||
}
|
||||
|
||||
{
|
||||
let gpu_pages = self.gpu_pages.read().await;
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
let metadata = serde_json::json!({
|
||||
"page_count": gpu_pages.len(),
|
||||
"timestamp": timestamp,
|
||||
"cache_type": "gpu"
|
||||
});
|
||||
|
||||
let metadata_path = cache_dir.join("gpu_pages_metadata.json");
|
||||
let metadata_str = serde_json::to_string_pretty(&metadata).map_err(|e| {
|
||||
InferenceError::kv_cache_error("persist", format!("Failed to serialize: {e}"))
|
||||
})?;
|
||||
fs::write(metadata_path, metadata_str).map_err(|e| {
|
||||
InferenceError::kv_cache_error("persist", format!("Failed to write: {e}"))
|
||||
})?;
|
||||
}
|
||||
|
||||
{
|
||||
let stats = self.stats.read().await;
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
let stats_json = serde_json::json!({
|
||||
"cache_hits": stats.cache_hits,
|
||||
"cache_misses": stats.cache_misses,
|
||||
"total_evictions": stats.evictions_lru + stats.evictions_lfu + stats.evictions_due_to_pressure,
|
||||
"timestamp": timestamp
|
||||
});
|
||||
|
||||
let stats_path = cache_dir.join("cache_stats.json");
|
||||
let stats_str = serde_json::to_string_pretty(&stats_json).map_err(|e| {
|
||||
InferenceError::kv_cache_error("persist", format!("Failed to serialize: {e}"))
|
||||
})?;
|
||||
fs::write(stats_path, stats_str).map_err(|e| {
|
||||
InferenceError::kv_cache_error("persist", format!("Failed to write: {e}"))
|
||||
})?;
|
||||
}
|
||||
|
||||
info!("Cache persistence completed to {:?}", cache_dir);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Recover cache from disk
|
||||
pub async fn recover_from_disk(&mut self) -> InferenceResult<()> {
|
||||
if !self.config.persistence_enabled {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
use std::fs;
|
||||
|
||||
let cache_dir = std::path::Path::new("./cache_persistence");
|
||||
if !cache_dir.exists() {
|
||||
info!("No cache persistence directory found, starting with empty cache");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let stats_path = cache_dir.join("cache_stats.json");
|
||||
if stats_path.exists() {
|
||||
match fs::read_to_string(&stats_path) {
|
||||
Ok(stats_json) => match serde_json::from_str::<serde_json::Value>(&stats_json) {
|
||||
Ok(stats_data) => {
|
||||
let mut stats = self.stats.write().await;
|
||||
if let Some(hits) = stats_data["cache_hits"].as_u64() {
|
||||
stats.cache_hits = hits as usize;
|
||||
}
|
||||
if let Some(misses) = stats_data["cache_misses"].as_u64() {
|
||||
stats.cache_misses = misses as usize;
|
||||
}
|
||||
if let Some(evictions) = stats_data["total_evictions"].as_u64() {
|
||||
stats.evictions_due_to_pressure = evictions as usize;
|
||||
}
|
||||
info!(
|
||||
"Recovered cache statistics: {} hits, {} misses",
|
||||
stats.cache_hits, stats.cache_misses
|
||||
);
|
||||
}
|
||||
Err(e) => warn!("Failed to parse cache stats: {}", e),
|
||||
},
|
||||
Err(e) => warn!("Failed to read cache stats: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
info!("Cache recovery completed from {:?}", cache_dir);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user