use super::{
entry::ResidentIfdsCsrCacheEntry, stats::ResidentIfdsCsrCacheStats,
};
use crate::ifds_gpu::{
free_resident_prepared_ifds_csr, upload_prepared_ifds_csr_resident, IfdsResidentDispatch,
PreparedIfdsCsr, ResidentPreparedIfdsCsr,
};
use crate::resident_cache_identity::{
ResidentGraphCacheIdentity, ResidentGraphCacheMissEvidence, ResidentGraphCacheMissReason,
};
use crate::resident_cache_lru::ResidentCacheLru;
use std::collections::HashMap;
use vyre::ResidentGraphReuseTelemetry;
pub struct ResidentIfdsCsrCache<R> {
entries: HashMap<ResidentGraphCacheIdentity, ResidentIfdsCsrCacheEntry<R>>,
lru: ResidentCacheLru<ResidentGraphCacheIdentity>,
max_retained_bytes: Option<usize>,
retained_bytes: usize,
clock: u64,
hits: u64,
misses: u64,
resident_uploads: u64,
resident_graph_reuse: ResidentGraphReuseTelemetry,
evictions: u64,
last_miss_reason: Option<ResidentGraphCacheMissReason>,
}
impl<R> Default for ResidentIfdsCsrCache<R> {
fn default() -> Self {
Self {
entries: HashMap::new(),
lru: ResidentCacheLru::default(),
max_retained_bytes: None,
retained_bytes: 0,
clock: 0,
hits: 0,
misses: 0,
resident_uploads: 0,
resident_graph_reuse: ResidentGraphReuseTelemetry::default(),
evictions: 0,
last_miss_reason: None,
}
}
}
impl<R> ResidentIfdsCsrCache<R>
where
R: Clone,
{
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_max_retained_bytes(max_retained_bytes: usize) -> Self {
Self {
max_retained_bytes: Some(max_retained_bytes),
..Self::default()
}
}
pub fn get_or_upload<D>(
&mut self,
dispatch: &D,
prepared: &PreparedIfdsCsr,
) -> Result<&ResidentPreparedIfdsCsr<R>, String>
where
D: IfdsResidentDispatch<Resource = R>,
{
let backend_id = dispatch.resident_backend_id();
let backend_version = dispatch.resident_backend_version();
let layout_hash = prepared.stable_layout_hash();
let edge_count = prepared.shape().edge_count;
let frontier_words = u32::try_from(prepared.frontier_words()).map_err(|error| {
format!(
"resident IFDS CSR frontier word count does not fit u32: {error}. Fix: shard the IFDS problem before resident caching."
)
})?;
let identity = ResidentGraphCacheIdentity::ifds_csr(
backend_id,
backend_version,
layout_hash,
prepared.node_count(),
edge_count,
frontier_words,
);
if let Some(entry) = self.entries.get(&identity) {
let retained_bytes = entry.retained_bytes;
if entry.graph.node_count() != prepared.node_count()
|| entry.graph.edge_count() != edge_count
|| entry.graph.frontier_words() != prepared.frontier_words()
|| entry.graph.stable_layout_hash() != layout_hash
{
return Err(format!(
"resident IFDS CSR cache entry for backend `{backend_id}` version `{backend_version}` has layout_hash={} nodes={} edges={} frontier_words={} but requested layout_hash={} nodes={} edges={} frontier_words={}. Fix: rebuild the IFDS resident CSR cache before dispatch.",
entry.graph.stable_layout_hash(),
entry.graph.node_count(),
entry.graph.edge_count(),
entry.graph.frontier_words(),
layout_hash,
prepared.node_count(),
edge_count,
prepared.frontier_words()
));
}
self.last_miss_reason = None;
self.hits = self
.hits
.checked_add(1)
.ok_or_else(|| {
"resident IFDS CSR cache hit counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
let retained_bytes_u64 = u64::try_from(retained_bytes).map_err(|error| {
format!(
"resident IFDS CSR retained byte count does not fit u64 during hit accounting: {error}. Fix: shard the IFDS graph before resident reuse."
)
})?;
self.resident_graph_reuse
.record_warm_reuse(retained_bytes_u64)
.map_err(|error| error.to_string())?;
let last_seen = self.next_cache_tick()?;
self.record_lru(identity.clone(), last_seen)?;
let entry = self.entries.get_mut(&identity).ok_or_else(|| {
"resident IFDS CSR cache key disappeared after LRU refresh. Fix: rebuild the resident IFDS cache; this indicates internal cache state corruption."
.to_string()
})?;
entry.last_seen = last_seen;
return Ok(&entry.graph);
}
self.last_miss_reason = Some(self.miss_reason_for_identity(&identity));
self.misses = self
.misses
.checked_add(1)
.ok_or_else(|| {
"resident IFDS CSR cache miss counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
let retained_bytes = prepared.retained_graph_bytes();
self.evict_until_room(dispatch, retained_bytes)?;
let resident = upload_prepared_ifds_csr_resident(dispatch, prepared)?;
self.resident_uploads = self
.resident_uploads
.checked_add(1)
.ok_or_else(|| {
"resident IFDS CSR cache upload counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
self.resident_graph_reuse
.record_cold_upload(u64::try_from(retained_bytes).map_err(|_| {
"resident IFDS CSR retained byte count exceeded u64. Fix: shard the IFDS graph before resident upload."
.to_string()
})?)
.map_err(|error| error.to_string())?;
let last_seen = self.next_cache_tick()?;
self.entries.insert(
identity.clone(),
ResidentIfdsCsrCacheEntry {
retained_bytes,
last_seen,
graph: resident,
},
);
self.record_lru(identity.clone(), last_seen)?;
self.retained_bytes = self
.retained_bytes
.checked_add(retained_bytes)
.ok_or_else(|| {
"resident IFDS CSR cache retained byte accounting overflowed usize. Fix: reduce retained graph budget or rebuild the resident IFDS cache."
.to_string()
})?;
Ok(&self
.entries
.get(&identity)
.ok_or_else(|| {
"resident IFDS CSR cache key disappeared after upload. Fix: rebuild the resident IFDS cache; this indicates internal cache state corruption."
.to_string()
})?
.graph)
}
#[must_use]
pub fn len(&self) -> usize {
self.entries.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
#[must_use]
pub fn stats(&self) -> ResidentIfdsCsrCacheStats {
ResidentIfdsCsrCacheStats {
hits: self.hits,
misses: self.misses,
resident_uploads: self.resident_uploads,
resident_upload_bytes: self.resident_graph_reuse.upload_bytes,
resident_avoided_upload_bytes: self.resident_graph_reuse.avoided_upload_bytes,
evictions: self.evictions,
retained_bytes: self.retained_bytes,
entries: self.entries.len(),
}
}
#[must_use]
pub fn resident_identities(&self) -> Vec<ResidentGraphCacheIdentity> {
let mut identities = self.entries.keys().cloned().collect::<Vec<_>>();
identities.sort_by(ResidentGraphCacheIdentity::cmp_stable);
identities
}
#[must_use]
pub fn miss_reason_for_identity(
&self,
requested: &ResidentGraphCacheIdentity,
) -> ResidentGraphCacheMissReason {
ResidentGraphCacheMissReason::classify(ResidentGraphCacheMissEvidence::from_identities(
self.entries.keys(),
requested,
))
}
#[must_use]
pub const fn last_miss_reason(&self) -> Option<ResidentGraphCacheMissReason> {
self.last_miss_reason
}
fn evict_until_room<D>(&mut self, dispatch: &D, incoming_bytes: usize) -> Result<(), String>
where
D: IfdsResidentDispatch<Resource = R>,
{
let Some(max_retained_bytes) = self.max_retained_bytes else {
return Ok(());
};
while !self.entries.is_empty()
&& self
.retained_bytes
.checked_add(incoming_bytes)
.is_none_or(|total| total > max_retained_bytes)
{
let Some(key) = self.pop_lru_key() else {
return Ok(());
};
let entry = self
.entries
.remove(&key)
.ok_or_else(|| {
"resident IFDS CSR eviction key disappeared before removal. Fix: rebuild the resident IFDS cache; this indicates internal cache state corruption."
.to_string()
})?;
self.retained_bytes = self
.retained_bytes
.checked_sub(entry.retained_bytes)
.ok_or_else(|| {
"resident IFDS CSR cache retained byte accounting underflowed during eviction. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
self.evictions = self
.evictions
.checked_add(1)
.ok_or_else(|| {
"resident IFDS CSR cache eviction counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
free_resident_prepared_ifds_csr(dispatch, entry.graph)?;
}
Ok(())
}
pub fn free_all<D>(&mut self, dispatch: &D) -> Result<(), String>
where
D: IfdsResidentDispatch<Resource = R>,
{
let mut first_error = None;
for (_, entry) in self.entries.drain() {
self.retained_bytes = self
.retained_bytes
.checked_sub(entry.retained_bytes)
.ok_or_else(|| {
"resident IFDS CSR cache retained byte accounting underflowed while freeing all entries. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
if let Err(error) = free_resident_prepared_ifds_csr(dispatch, entry.graph) {
first_error.get_or_insert(error);
}
}
self.lru.clear();
match first_error {
Some(error) => Err(error),
None => Ok(()),
}
}
fn next_cache_tick(&mut self) -> Result<u64, String> {
self.clock = self
.clock
.checked_add(1)
.ok_or_else(|| {
"resident IFDS CSR cache logical clock overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
.to_string()
})?;
Ok(self.clock)
}
fn pop_lru_key(&mut self) -> Option<ResidentGraphCacheIdentity> {
self.lru.pop_valid(
|key| self.entries.get(key).map(|entry| entry.last_seen),
|| self.entries.keys().next().cloned(),
)
}
fn record_lru(
&mut self,
key: ResidentGraphCacheIdentity,
last_seen: u64,
) -> Result<(), String> {
self.lru.record(
key,
last_seen,
self.entries
.iter()
.map(|(key, entry)| (key.clone(), entry.last_seen)),
"resident IFDS CSR cache",
)
}
#[cfg(test)]
pub(super) fn lru_len_for_tests(&self) -> usize {
self.lru.len()
}
}