weir/ifds_resident_cache/
cache.rs

1use super::{
2    entry::ResidentIfdsCsrCacheEntry, stats::ResidentIfdsCsrCacheStats,
3};
4use crate::ifds_gpu::{
5    free_resident_prepared_ifds_csr, upload_prepared_ifds_csr_resident, IfdsResidentDispatch,
6    PreparedIfdsCsr, ResidentPreparedIfdsCsr,
7};
8use crate::resident_cache_identity::{
9    ResidentGraphCacheIdentity, ResidentGraphCacheMissEvidence, ResidentGraphCacheMissReason,
10};
11use crate::resident_cache_lru::ResidentCacheLru;
12use std::collections::HashMap;
13use vyre::ResidentGraphReuseTelemetry;
14
15/// Backend-resident IFDS CSR cache keyed by normalized prepared layout.
16pub struct ResidentIfdsCsrCache<R> {
17    entries: HashMap<ResidentGraphCacheIdentity, ResidentIfdsCsrCacheEntry<R>>,
18    lru: ResidentCacheLru<ResidentGraphCacheIdentity>,
19    max_retained_bytes: Option<usize>,
20    retained_bytes: usize,
21    clock: u64,
22    hits: u64,
23    misses: u64,
24    resident_uploads: u64,
25    resident_graph_reuse: ResidentGraphReuseTelemetry,
26    evictions: u64,
27    last_miss_reason: Option<ResidentGraphCacheMissReason>,
28}
29
30impl<R> Default for ResidentIfdsCsrCache<R> {
31    fn default() -> Self {
32        Self {
33            entries: HashMap::new(),
34            lru: ResidentCacheLru::default(),
35            max_retained_bytes: None,
36            retained_bytes: 0,
37            clock: 0,
38            hits: 0,
39            misses: 0,
40            resident_uploads: 0,
41            resident_graph_reuse: ResidentGraphReuseTelemetry::default(),
42            evictions: 0,
43            last_miss_reason: None,
44        }
45    }
46}
47
48impl<R> ResidentIfdsCsrCache<R>
49where
50    R: Clone,
51{
52    /// Create an empty resident IFDS CSR cache.
53    #[must_use]
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    /// Create a cache with a soft upper bound on retained resident graph bytes.
59    ///
60    /// The newest entry is always retained so one oversized graph may exceed the
61    /// limit; older entries are evicted before uploading a cache miss.
62    #[must_use]
63    pub fn with_max_retained_bytes(max_retained_bytes: usize) -> Self {
64        Self {
65            max_retained_bytes: Some(max_retained_bytes),
66            ..Self::default()
67        }
68    }
69
70    /// Return a resident IFDS CSR for `prepared`, uploading on cache miss.
71    pub fn get_or_upload<D>(
72        &mut self,
73        dispatch: &D,
74        prepared: &PreparedIfdsCsr,
75    ) -> Result<&ResidentPreparedIfdsCsr<R>, String>
76    where
77        D: IfdsResidentDispatch<Resource = R>,
78    {
79        let backend_id = dispatch.resident_backend_id();
80        let backend_version = dispatch.resident_backend_version();
81        let layout_hash = prepared.stable_layout_hash();
82        let edge_count = prepared.shape().edge_count;
83        let frontier_words = u32::try_from(prepared.frontier_words()).map_err(|error| {
84            format!(
85                "resident IFDS CSR frontier word count does not fit u32: {error}. Fix: shard the IFDS problem before resident caching."
86            )
87        })?;
88        let identity = ResidentGraphCacheIdentity::ifds_csr(
89            backend_id,
90            backend_version,
91            layout_hash,
92            prepared.node_count(),
93            edge_count,
94            frontier_words,
95        );
96        if let Some(entry) = self.entries.get(&identity) {
97            let retained_bytes = entry.retained_bytes;
98            if entry.graph.node_count() != prepared.node_count()
99                || entry.graph.edge_count() != edge_count
100                || entry.graph.frontier_words() != prepared.frontier_words()
101                || entry.graph.stable_layout_hash() != layout_hash
102            {
103                return Err(format!(
104                    "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.",
105                    entry.graph.stable_layout_hash(),
106                    entry.graph.node_count(),
107                    entry.graph.edge_count(),
108                    entry.graph.frontier_words(),
109                    layout_hash,
110                    prepared.node_count(),
111                    edge_count,
112                    prepared.frontier_words()
113                ));
114            }
115            self.last_miss_reason = None;
116            self.hits = self
117                .hits
118                .checked_add(1)
119                .ok_or_else(|| {
120                    "resident IFDS CSR cache hit counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
121                        .to_string()
122                })?;
123            let retained_bytes_u64 = u64::try_from(retained_bytes).map_err(|error| {
124                format!(
125                    "resident IFDS CSR retained byte count does not fit u64 during hit accounting: {error}. Fix: shard the IFDS graph before resident reuse."
126                )
127            })?;
128            self.resident_graph_reuse
129                .record_warm_reuse(retained_bytes_u64)
130                .map_err(|error| error.to_string())?;
131            let last_seen = self.next_cache_tick()?;
132            self.record_lru(identity.clone(), last_seen)?;
133            let entry = self.entries.get_mut(&identity).ok_or_else(|| {
134                "resident IFDS CSR cache key disappeared after LRU refresh. Fix: rebuild the resident IFDS cache; this indicates internal cache state corruption."
135                    .to_string()
136            })?;
137            entry.last_seen = last_seen;
138            return Ok(&entry.graph);
139        }
140        self.last_miss_reason = Some(self.miss_reason_for_identity(&identity));
141        self.misses = self
142            .misses
143            .checked_add(1)
144            .ok_or_else(|| {
145                "resident IFDS CSR cache miss counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
146                    .to_string()
147            })?;
148        let retained_bytes = prepared.retained_graph_bytes();
149        self.evict_until_room(dispatch, retained_bytes)?;
150        let resident = upload_prepared_ifds_csr_resident(dispatch, prepared)?;
151        self.resident_uploads = self
152            .resident_uploads
153            .checked_add(1)
154            .ok_or_else(|| {
155                "resident IFDS CSR cache upload counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
156                    .to_string()
157            })?;
158        self.resident_graph_reuse
159            .record_cold_upload(u64::try_from(retained_bytes).map_err(|_| {
160                "resident IFDS CSR retained byte count exceeded u64. Fix: shard the IFDS graph before resident upload."
161                    .to_string()
162            })?)
163            .map_err(|error| error.to_string())?;
164        let last_seen = self.next_cache_tick()?;
165        self.entries.insert(
166            identity.clone(),
167            ResidentIfdsCsrCacheEntry {
168                retained_bytes,
169                last_seen,
170                graph: resident,
171            },
172        );
173        self.record_lru(identity.clone(), last_seen)?;
174        self.retained_bytes = self
175            .retained_bytes
176            .checked_add(retained_bytes)
177            .ok_or_else(|| {
178                "resident IFDS CSR cache retained byte accounting overflowed usize. Fix: reduce retained graph budget or rebuild the resident IFDS cache."
179                    .to_string()
180            })?;
181        Ok(&self
182            .entries
183            .get(&identity)
184            .ok_or_else(|| {
185                "resident IFDS CSR cache key disappeared after upload. Fix: rebuild the resident IFDS cache; this indicates internal cache state corruption."
186                    .to_string()
187            })?
188            .graph)
189    }
190
191    /// Number of resident IFDS CSR entries retained by the cache.
192    #[must_use]
193    pub fn len(&self) -> usize {
194        self.entries.len()
195    }
196
197    /// Whether the cache is empty.
198    #[must_use]
199    pub fn is_empty(&self) -> bool {
200        self.entries.is_empty()
201    }
202
203    /// Return point-in-time cache counters.
204    #[must_use]
205    pub fn stats(&self) -> ResidentIfdsCsrCacheStats {
206        ResidentIfdsCsrCacheStats {
207            hits: self.hits,
208            misses: self.misses,
209            resident_uploads: self.resident_uploads,
210            resident_upload_bytes: self.resident_graph_reuse.upload_bytes,
211            resident_avoided_upload_bytes: self.resident_graph_reuse.avoided_upload_bytes,
212            evictions: self.evictions,
213            retained_bytes: self.retained_bytes,
214            entries: self.entries.len(),
215        }
216    }
217
218    /// Deterministic cache identity snapshot for diagnostics and release
219    /// evidence.
220    #[must_use]
221    pub fn resident_identities(&self) -> Vec<ResidentGraphCacheIdentity> {
222        let mut identities = self.entries.keys().cloned().collect::<Vec<_>>();
223        identities.sort_by(ResidentGraphCacheIdentity::cmp_stable);
224        identities
225    }
226
227    /// Explain why `requested` would miss the current resident IFDS cache.
228    #[must_use]
229    pub fn miss_reason_for_identity(
230        &self,
231        requested: &ResidentGraphCacheIdentity,
232    ) -> ResidentGraphCacheMissReason {
233        ResidentGraphCacheMissReason::classify(ResidentGraphCacheMissEvidence::from_identities(
234            self.entries.keys(),
235            requested,
236        ))
237    }
238
239    /// Reason attached to the most recent cache miss. Hits clear this field.
240    #[must_use]
241    pub const fn last_miss_reason(&self) -> Option<ResidentGraphCacheMissReason> {
242        self.last_miss_reason
243    }
244
245    fn evict_until_room<D>(&mut self, dispatch: &D, incoming_bytes: usize) -> Result<(), String>
246    where
247        D: IfdsResidentDispatch<Resource = R>,
248    {
249        let Some(max_retained_bytes) = self.max_retained_bytes else {
250            return Ok(());
251        };
252        while !self.entries.is_empty()
253            && self
254                .retained_bytes
255                .checked_add(incoming_bytes)
256                .is_none_or(|total| total > max_retained_bytes)
257        {
258            let Some(key) = self.pop_lru_key() else {
259                return Ok(());
260            };
261            let entry = self
262                .entries
263                .remove(&key)
264                .ok_or_else(|| {
265                    "resident IFDS CSR eviction key disappeared before removal. Fix: rebuild the resident IFDS cache; this indicates internal cache state corruption."
266                        .to_string()
267                })?;
268            self.retained_bytes = self
269                .retained_bytes
270                .checked_sub(entry.retained_bytes)
271                .ok_or_else(|| {
272                    "resident IFDS CSR cache retained byte accounting underflowed during eviction. Fix: rebuild the resident IFDS cache before reuse."
273                        .to_string()
274                })?;
275            self.evictions = self
276                .evictions
277                .checked_add(1)
278                .ok_or_else(|| {
279                    "resident IFDS CSR cache eviction counter overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
280                        .to_string()
281                })?;
282            free_resident_prepared_ifds_csr(dispatch, entry.graph)?;
283        }
284        Ok(())
285    }
286
287    /// Free every resident IFDS CSR owned by this cache.
288    pub fn free_all<D>(&mut self, dispatch: &D) -> Result<(), String>
289    where
290        D: IfdsResidentDispatch<Resource = R>,
291    {
292        let mut first_error = None;
293        for (_, entry) in self.entries.drain() {
294            self.retained_bytes = self
295                .retained_bytes
296                .checked_sub(entry.retained_bytes)
297                .ok_or_else(|| {
298                    "resident IFDS CSR cache retained byte accounting underflowed while freeing all entries. Fix: rebuild the resident IFDS cache before reuse."
299                        .to_string()
300                })?;
301            if let Err(error) = free_resident_prepared_ifds_csr(dispatch, entry.graph) {
302                first_error.get_or_insert(error);
303            }
304        }
305        self.lru.clear();
306        match first_error {
307            Some(error) => Err(error),
308            None => Ok(()),
309        }
310    }
311
312    fn next_cache_tick(&mut self) -> Result<u64, String> {
313        self.clock = self
314            .clock
315            .checked_add(1)
316            .ok_or_else(|| {
317                "resident IFDS CSR cache logical clock overflowed u64. Fix: rebuild the resident IFDS cache before reuse."
318                    .to_string()
319            })?;
320        Ok(self.clock)
321    }
322
323    fn pop_lru_key(&mut self) -> Option<ResidentGraphCacheIdentity> {
324        self.lru.pop_valid(
325            |key| self.entries.get(key).map(|entry| entry.last_seen),
326            || self.entries.keys().next().cloned(),
327        )
328    }
329
330    fn record_lru(
331        &mut self,
332        key: ResidentGraphCacheIdentity,
333        last_seen: u64,
334    ) -> Result<(), String> {
335        self.lru.record(
336            key,
337            last_seen,
338            self.entries
339                .iter()
340                .map(|(key, entry)| (key.clone(), entry.last_seen)),
341            "resident IFDS CSR cache",
342        )
343    }
344
345    #[cfg(test)]
346    pub(super) fn lru_len_for_tests(&self) -> usize {
347        self.lru.len()
348    }
349}
weir/ifds_resident_cache/cache.rs

weir/ifds_resident_cache/
cache.rs
