eredu_runtime/
cache.rs

1//! Backend-neutral ownership and admission algorithms for live model caches.
2//!
3//! A cache backend owns concrete tensors, transfer objects, and persistent
4//! files. This module owns both each session's block/lease/tail lifecycle and
5//! the process-wide accounting boundary shared by independently managed cache
6//! sessions. Reservations are RAII tokens: an admitted resource remains charged
7//! until its exact backend transition either publishes the bytes into a
8//! registered manager or drops the reservation.
9
10mod executor;
11mod lifecycle;
12mod persistence;
13mod policy;
14mod storage;
15mod telemetry;
16mod worker;
17
18pub use executor::{
19    CacheIoAdmission, CacheIoCompletionDisposition, CacheIoExecutionState,
20    CacheIoExecutionStateError, CacheIoPreparation, CacheIoStartDisposition,
21};
22pub use lifecycle::{CacheBlockLifecycle, CacheLifecycleError, MutableCacheTail};
23pub use persistence::{
24    finalize_prompt_cache_shard, hash_prompt_cache_shard_payload, inspect_prompt_cache,
25    prompt_cache_rank_path, resolve_prompt_cache_root, safe_prompt_cache_shard_path,
26    validate_prompt_cache_manifest, LiveCacheBlockPublication, LiveCachePublicationError,
27    PromptCachePersistenceError, PromptCachePublication, ReversiblePromptCachePublication,
28    MAX_PROMPT_CACHE_SHARD_HEADER_BYTES, PROMPT_CACHE_CURRENT_FILE,
29    PROMPT_CACHE_GENERATIONS_DIRECTORY,
30};
31pub use policy::{
32    CacheResidencyConfigurationError, CacheResidencyPolicy, LiveCacheDiskPolicy, PagedCacheOptions,
33};
34pub use storage::{
35    CacheBlockStorage, CacheHostDemotionOperation, CacheHostPromotion, CacheIoOperation,
36    CacheIoOperationKey, CacheIoOperationKind, CacheStorageError, CacheStoragePhase,
37};
38pub use telemetry::{
39    CacheLayerResidencyReport, CacheLayerResidencyStats, CacheResidencyReport,
40    CacheResidencyTelemetry, CACHE_RESIDENCY_LAYER_REPORT_LIMIT,
41};
42pub use worker::{
43    CacheIoSubmission, CacheIoSubmissionOutcome, CacheIoTicket, CacheIoWorker, CacheIoWorkerError,
44};
45
46use serde::{Deserialize, Serialize};
47use std::{
48    collections::BTreeMap,
49    sync::{
50        atomic::{AtomicU64, Ordering},
51        Arc, Mutex,
52    },
53};
54
55static NEXT_CACHE_POOL_ID: AtomicU64 = AtomicU64::new(1);
56static NEXT_CACHE_POOL_RESERVATION_ID: AtomicU64 = AtomicU64::new(1);
57
58/// Process-wide finite limits shared by independently owned live caches.
59#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
60pub struct CachePoolLimits {
61    device_bytes: u64,
62    host_bytes: u64,
63    transfer_in_flight_bytes: u64,
64    disk_bytes: u64,
65}
66
67impl CachePoolLimits {
68    /// Creates finite aggregate cache limits.
69    ///
70    /// Device and transfer capacity must be nonzero. Zero host or disk
71    /// capacity explicitly disables that tier.
72    pub fn new(
73        device_bytes: u64,
74        host_bytes: u64,
75        transfer_in_flight_bytes: u64,
76        disk_bytes: u64,
77    ) -> Result<Self, CachePoolError> {
78        if device_bytes == 0 {
79            return Err(CachePoolError::InvalidLimits(
80                "cache pool device budget must be nonzero",
81            ));
82        }
83        if transfer_in_flight_bytes == 0 {
84            return Err(CachePoolError::InvalidLimits(
85                "cache pool transfer-in-flight budget must be nonzero",
86            ));
87        }
88        Ok(Self {
89            device_bytes,
90            host_bytes,
91            transfer_in_flight_bytes,
92            disk_bytes,
93        })
94    }
95
96    /// Aggregate execution-device cache capacity.
97    pub const fn device_bytes(self) -> u64 {
98        self.device_bytes
99    }
100
101    /// Aggregate physical host-allocation capacity.
102    pub const fn host_bytes(self) -> u64 {
103        self.host_bytes
104    }
105
106    /// Aggregate bytes retained until exact transfer completion.
107    pub const fn transfer_in_flight_bytes(self) -> u64 {
108        self.transfer_in_flight_bytes
109    }
110
111    /// Aggregate live-cache disk capacity.
112    pub const fn disk_bytes(self) -> u64 {
113        self.disk_bytes
114    }
115}
116
117/// Occupancy on each independently admitted cache resource axis.
118#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
119pub struct CachePoolUsage {
120    /// Concrete execution-device allocations.
121    pub device_bytes: u64,
122    /// Concrete physical host allocations.
123    pub host_bytes: u64,
124    /// Resources retained by submitted transfers.
125    pub transfer_in_flight_bytes: u64,
126    /// Live-cache files plus pending file reservations.
127    pub disk_bytes: u64,
128}
129
130impl CachePoolUsage {
131    fn checked_add(self, other: Self) -> Option<Self> {
132        Some(Self {
133            device_bytes: self.device_bytes.checked_add(other.device_bytes)?,
134            host_bytes: self.host_bytes.checked_add(other.host_bytes)?,
135            transfer_in_flight_bytes: self
136                .transfer_in_flight_bytes
137                .checked_add(other.transfer_in_flight_bytes)?,
138            disk_bytes: self.disk_bytes.checked_add(other.disk_bytes)?,
139        })
140    }
141
142    fn checked_sub(self, other: Self) -> Option<Self> {
143        Some(Self {
144            device_bytes: self.device_bytes.checked_sub(other.device_bytes)?,
145            host_bytes: self.host_bytes.checked_sub(other.host_bytes)?,
146            transfer_in_flight_bytes: self
147                .transfer_in_flight_bytes
148                .checked_sub(other.transfer_in_flight_bytes)?,
149            disk_bytes: self.disk_bytes.checked_sub(other.disk_bytes)?,
150        })
151    }
152}
153
154/// Resource axis named by an aggregate admission failure.
155#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
156#[serde(rename_all = "snake_case")]
157pub enum CachePoolResource {
158    /// Execution-device cache allocations.
159    Device,
160    /// Physical host cache allocations.
161    Host,
162    /// Resources retained until exact transfer completion.
163    TransferInFlight,
164    /// Live-cache disk files and pending writes.
165    Disk,
166}
167
168/// Aggregate process-pool occupancy and high-water marks.
169#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
170pub struct CachePoolReport {
171    /// Stable process-local pool identity.
172    pub pool_id: u64,
173    /// Number of registered live cache managers.
174    pub managers: usize,
175    /// Current aggregate device bytes.
176    pub current_device_bytes: u64,
177    /// Peak aggregate device bytes successfully admitted.
178    pub peak_device_bytes: u64,
179    /// Current aggregate physical host capacity.
180    pub current_host_bytes: u64,
181    /// Peak aggregate physical host capacity successfully admitted.
182    pub peak_host_bytes: u64,
183    /// Current bytes retained by in-flight transfers.
184    pub current_transfer_in_flight_bytes: u64,
185    /// Peak bytes retained by in-flight transfers.
186    pub peak_transfer_in_flight_bytes: u64,
187    /// Current disk bytes, including pending write reservations.
188    pub current_disk_bytes: u64,
189    /// Peak disk bytes successfully admitted.
190    pub peak_disk_bytes: u64,
191    /// Aggregate finite limits.
192    pub limits: CachePoolLimits,
193}
194
195#[derive(Debug)]
196struct CachePoolState {
197    managers: BTreeMap<u64, CachePoolUsage>,
198    reservations: BTreeMap<u64, CachePoolUsage>,
199    current: CachePoolUsage,
200    peak: CachePoolUsage,
201}
202
203/// Shareable aggregate ownership boundary for scheduler and standalone caches.
204#[derive(Clone)]
205pub struct CacheResidencyPool {
206    id: u64,
207    limits: CachePoolLimits,
208    state: Arc<Mutex<CachePoolState>>,
209}
210
211impl std::fmt::Debug for CacheResidencyPool {
212    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
213        formatter
214            .debug_struct("CacheResidencyPool")
215            .field("id", &self.id)
216            .field("limits", &self.limits)
217            .finish_non_exhaustive()
218    }
219}
220
221impl PartialEq for CacheResidencyPool {
222    fn eq(&self, other: &Self) -> bool {
223        Arc::ptr_eq(&self.state, &other.state)
224    }
225}
226
227impl Eq for CacheResidencyPool {}
228
229impl CacheResidencyPool {
230    /// Creates an empty process pool under aggregate finite limits.
231    pub fn new(limits: CachePoolLimits) -> Self {
232        Self {
233            id: NEXT_CACHE_POOL_ID.fetch_add(1, Ordering::Relaxed),
234            limits,
235            state: Arc::new(Mutex::new(CachePoolState {
236                managers: BTreeMap::new(),
237                reservations: BTreeMap::new(),
238                current: CachePoolUsage::default(),
239                peak: CachePoolUsage::default(),
240            })),
241        }
242    }
243
244    /// Returns this pool's stable process-local identity.
245    pub const fn id(&self) -> u64 {
246        self.id
247    }
248
249    /// Returns aggregate finite limits.
250    pub const fn limits(&self) -> CachePoolLimits {
251        self.limits
252    }
253
254    /// Registers a manager and returns its exact RAII membership token.
255    pub fn register_manager(&self, manager: u64) -> Result<CachePoolMembership, CachePoolError> {
256        let mut state = self.state.lock().map_err(|_| CachePoolError::Poisoned)?;
257        if state.managers.contains_key(&manager) {
258            return Err(CachePoolError::DuplicateManager { manager });
259        }
260        state.managers.insert(manager, CachePoolUsage::default());
261        Ok(CachePoolMembership {
262            manager,
263            pool: self.clone(),
264        })
265    }
266
267    /// Replaces the published occupancy owned by one registered manager.
268    ///
269    /// Backends use reservations to admit growth before allocating it, publish
270    /// the resulting concrete occupancy here, and then release the reservation.
271    /// Publication records physical truth even when an unavoidable allocation
272    /// boundary temporarily exceeds a limit; subsequent admissions fail and
273    /// the backend must synchronously rebalance or roll back that allocation.
274    pub fn update_manager(
275        &self,
276        manager: u64,
277        usage: CachePoolUsage,
278    ) -> Result<CachePoolUsage, CachePoolError> {
279        let mut state = self.state.lock().map_err(|_| CachePoolError::Poisoned)?;
280        let old = *state
281            .managers
282            .get(&manager)
283            .ok_or(CachePoolError::UnknownManager { manager })?;
284        let current = state
285            .current
286            .checked_sub(old)
287            .and_then(|current| current.checked_add(usage))
288            .ok_or(CachePoolError::AccountingOverflow {
289                operation: "manager occupancy publication",
290            })?;
291        state.managers.insert(manager, usage);
292        state.current = current;
293        update_peaks(&mut state, self.limits);
294        Ok(state.current)
295    }
296
297    /// Atomically admits additional occupancy until its RAII token is dropped.
298    pub fn reserve(&self, usage: CachePoolUsage) -> Result<CachePoolReservation, CachePoolError> {
299        let mut state = self.state.lock().map_err(|_| CachePoolError::Poisoned)?;
300        validate_additional(state.current, usage, self.limits)?;
301        let required =
302            state
303                .current
304                .checked_add(usage)
305                .ok_or(CachePoolError::AccountingOverflow {
306                    operation: "temporary admission",
307                })?;
308        let reservation = NEXT_CACHE_POOL_RESERVATION_ID.fetch_add(1, Ordering::Relaxed);
309        state.reservations.insert(reservation, usage);
310        state.current = required;
311        update_peaks(&mut state, self.limits);
312        Ok(CachePoolReservation {
313            reservation,
314            pool: self.clone(),
315        })
316    }
317
318    /// Atomically admits resources retained by one submitted transfer.
319    pub fn reserve_transfer(&self, bytes: u64) -> Result<CachePoolReservation, CachePoolError> {
320        self.reserve(CachePoolUsage {
321            transfer_in_flight_bytes: bytes,
322            ..CachePoolUsage::default()
323        })
324    }
325
326    /// Returns aggregate current occupancy and high-water marks.
327    pub fn report(&self) -> Result<CachePoolReport, CachePoolError> {
328        let state = self.state.lock().map_err(|_| CachePoolError::Poisoned)?;
329        Ok(CachePoolReport {
330            pool_id: self.id,
331            managers: state.managers.len(),
332            current_device_bytes: state.current.device_bytes,
333            peak_device_bytes: state.peak.device_bytes,
334            current_host_bytes: state.current.host_bytes,
335            peak_host_bytes: state.peak.host_bytes,
336            current_transfer_in_flight_bytes: state.current.transfer_in_flight_bytes,
337            peak_transfer_in_flight_bytes: state.peak.transfer_in_flight_bytes,
338            current_disk_bytes: state.current.disk_bytes,
339            peak_disk_bytes: state.peak.disk_bytes,
340            limits: self.limits,
341        })
342    }
343
344    fn remove_manager(&self, manager: u64) {
345        if let Ok(mut state) = self.state.lock() {
346            if let Some(previous) = state.managers.get(&manager).copied() {
347                if let Some(current) = state.current.checked_sub(previous) {
348                    state.managers.remove(&manager);
349                    state.current = current;
350                }
351            }
352        }
353    }
354}
355
356/// Exact ownership of a temporary aggregate cache admission.
357#[derive(Debug)]
358pub struct CachePoolReservation {
359    reservation: u64,
360    pool: CacheResidencyPool,
361}
362
363impl Drop for CachePoolReservation {
364    fn drop(&mut self) {
365        if let Ok(mut state) = self.pool.state.lock() {
366            if let Some(usage) = state.reservations.get(&self.reservation).copied() {
367                if let Some(current) = state.current.checked_sub(usage) {
368                    state.reservations.remove(&self.reservation);
369                    state.current = current;
370                }
371            }
372        }
373    }
374}
375
376/// Exact ownership of one manager's aggregate pool contribution.
377#[derive(Debug)]
378pub struct CachePoolMembership {
379    manager: u64,
380    pool: CacheResidencyPool,
381}
382
383impl CachePoolMembership {
384    /// Pool to which this manager is registered.
385    pub const fn pool(&self) -> &CacheResidencyPool {
386        &self.pool
387    }
388}
389
390impl Drop for CachePoolMembership {
391    fn drop(&mut self) {
392        self.pool.remove_manager(self.manager);
393    }
394}
395
396fn validate_additional(
397    current: CachePoolUsage,
398    usage: CachePoolUsage,
399    limits: CachePoolLimits,
400) -> Result<(), CachePoolError> {
401    let required = current
402        .checked_add(usage)
403        .ok_or(CachePoolError::AccountingOverflow {
404            operation: "admission validation",
405        })?;
406    for (resource, additional, required, budget) in [
407        (
408            CachePoolResource::Device,
409            usage.device_bytes,
410            required.device_bytes,
411            limits.device_bytes,
412        ),
413        (
414            CachePoolResource::Host,
415            usage.host_bytes,
416            required.host_bytes,
417            limits.host_bytes,
418        ),
419        (
420            CachePoolResource::TransferInFlight,
421            usage.transfer_in_flight_bytes,
422            required.transfer_in_flight_bytes,
423            limits.transfer_in_flight_bytes,
424        ),
425        (
426            CachePoolResource::Disk,
427            usage.disk_bytes,
428            required.disk_bytes,
429            limits.disk_bytes,
430        ),
431    ] {
432        if additional != 0 && required > budget {
433            return Err(CachePoolError::BudgetExceeded {
434                resource,
435                required,
436                budget,
437            });
438        }
439    }
440    Ok(())
441}
442
443fn update_peaks(state: &mut CachePoolState, limits: CachePoolLimits) {
444    if state.current.device_bytes <= limits.device_bytes {
445        state.peak.device_bytes = state.peak.device_bytes.max(state.current.device_bytes);
446    }
447    if state.current.host_bytes <= limits.host_bytes {
448        state.peak.host_bytes = state.peak.host_bytes.max(state.current.host_bytes);
449    }
450    if state.current.transfer_in_flight_bytes <= limits.transfer_in_flight_bytes {
451        state.peak.transfer_in_flight_bytes = state
452            .peak
453            .transfer_in_flight_bytes
454            .max(state.current.transfer_in_flight_bytes);
455    }
456    if state.current.disk_bytes <= limits.disk_bytes {
457        state.peak.disk_bytes = state.peak.disk_bytes.max(state.current.disk_bytes);
458    }
459}
460
461/// Backend-neutral aggregate cache ownership failure.
462#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
463pub enum CachePoolError {
464    /// Aggregate limits had no device or transfer capacity.
465    #[error("invalid cache pool limits: {0}")]
466    InvalidLimits(&'static str),
467    /// A manager identity was registered more than once.
468    #[error("cache pool manager {manager} is already registered")]
469    DuplicateManager {
470        /// Duplicate manager identity.
471        manager: u64,
472    },
473    /// Occupancy was published for an unregistered manager.
474    #[error("cache pool manager {manager} is not registered")]
475    UnknownManager {
476        /// Missing manager identity.
477        manager: u64,
478    },
479    /// Aggregate cache accounting exceeded one finite resource limit.
480    #[error(
481        "cache pool {resource:?} budget exceeded: required {required} bytes, budget {budget} bytes"
482    )]
483    BudgetExceeded {
484        /// Exhausted resource axis.
485        resource: CachePoolResource,
486        /// Aggregate bytes required.
487        required: u64,
488        /// Finite pool budget.
489        budget: u64,
490    },
491    /// Checked ownership accounting overflowed.
492    #[error("cache pool accounting overflow during {operation}")]
493    AccountingOverflow {
494        /// Stable accounting transition.
495        operation: &'static str,
496    },
497    /// Shared ownership state was poisoned by a panic.
498    #[error("cache residency pool state is poisoned")]
499    Poisoned,
500}
501
502#[cfg(test)]
503mod tests {
504    use super::*;
505    use std::sync::{mpsc, Barrier};
506
507    fn pool() -> CacheResidencyPool {
508        CacheResidencyPool::new(CachePoolLimits::new(16, 12, 10, 8).unwrap())
509    }
510
511    #[test]
512    fn manager_membership_owns_published_occupancy() {
513        let pool = pool();
514        let membership = pool.register_manager(7).unwrap();
515        pool.update_manager(
516            7,
517            CachePoolUsage {
518                device_bytes: 8,
519                host_bytes: 4,
520                ..CachePoolUsage::default()
521            },
522        )
523        .unwrap();
524        let report = pool.report().unwrap();
525        assert_eq!(report.managers, 1);
526        assert_eq!(report.current_device_bytes, 8);
527        assert_eq!(report.current_host_bytes, 4);
528
529        drop(membership);
530        let report = pool.report().unwrap();
531        assert_eq!(report.managers, 0);
532        assert_eq!(report.current_device_bytes, 0);
533        assert_eq!(report.current_host_bytes, 0);
534    }
535
536    #[test]
537    fn reservation_is_atomic_and_released_by_exact_owner() {
538        let pool = pool();
539        let first = pool.reserve_transfer(6).unwrap();
540        assert_eq!(pool.report().unwrap().current_transfer_in_flight_bytes, 6);
541        assert_eq!(
542            pool.reserve_transfer(5).unwrap_err(),
543            CachePoolError::BudgetExceeded {
544                resource: CachePoolResource::TransferInFlight,
545                required: 11,
546                budget: 10,
547            }
548        );
549        assert_eq!(pool.report().unwrap().current_transfer_in_flight_bytes, 6);
550        drop(first);
551        assert_eq!(pool.report().unwrap().current_transfer_in_flight_bytes, 0);
552    }
553
554    #[test]
555    fn independent_resource_axes_fail_closed() {
556        let pool = pool();
557        for (usage, resource, required, budget) in [
558            (
559                CachePoolUsage {
560                    device_bytes: 17,
561                    ..CachePoolUsage::default()
562                },
563                CachePoolResource::Device,
564                17,
565                16,
566            ),
567            (
568                CachePoolUsage {
569                    host_bytes: 13,
570                    ..CachePoolUsage::default()
571                },
572                CachePoolResource::Host,
573                13,
574                12,
575            ),
576            (
577                CachePoolUsage {
578                    disk_bytes: 9,
579                    ..CachePoolUsage::default()
580                },
581                CachePoolResource::Disk,
582                9,
583                8,
584            ),
585        ] {
586            assert_eq!(
587                pool.reserve(usage).unwrap_err(),
588                CachePoolError::BudgetExceeded {
589                    resource,
590                    required,
591                    budget,
592                }
593            );
594        }
595        assert_eq!(pool.report().unwrap().current_device_bytes, 0);
596        assert_eq!(pool.report().unwrap().current_host_bytes, 0);
597        assert_eq!(pool.report().unwrap().current_disk_bytes, 0);
598    }
599
600    #[test]
601    fn reports_and_limits_round_trip_without_backend_types() {
602        let pool = pool();
603        let _membership = pool.register_manager(1).unwrap();
604        pool.update_manager(
605            1,
606            CachePoolUsage {
607                device_bytes: 8,
608                disk_bytes: 4,
609                ..CachePoolUsage::default()
610            },
611        )
612        .unwrap();
613        let report = pool.report().unwrap();
614        let encoded = serde_json::to_string(&report).unwrap();
615        assert_eq!(
616            serde_json::from_str::<CachePoolReport>(&encoded).unwrap(),
617            report
618        );
619    }
620
621    #[test]
622    fn concurrent_admission_has_one_atomic_winner() {
623        let pool = CacheResidencyPool::new(CachePoolLimits::new(64, 10, 64, 0).unwrap());
624        let start = Arc::new(Barrier::new(3));
625        let finish = Arc::new(Barrier::new(3));
626        let (sender, receiver) = mpsc::channel();
627        let handles = (0..2)
628            .map(|_| {
629                let pool = pool.clone();
630                let start = Arc::clone(&start);
631                let finish = Arc::clone(&finish);
632                let sender = sender.clone();
633                std::thread::spawn(move || {
634                    start.wait();
635                    let admission = pool.reserve(CachePoolUsage {
636                        host_bytes: 6,
637                        ..CachePoolUsage::default()
638                    });
639                    sender.send(admission.is_ok()).unwrap();
640                    finish.wait();
641                    drop(admission);
642                })
643            })
644            .collect::<Vec<_>>();
645        drop(sender);
646        start.wait();
647        let admitted = [receiver.recv().unwrap(), receiver.recv().unwrap()];
648        assert_eq!(admitted.into_iter().filter(|value| *value).count(), 1);
649        assert_eq!(pool.report().unwrap().current_host_bytes, 6);
650        finish.wait();
651        for handle in handles {
652            handle.join().unwrap();
653        }
654        assert_eq!(pool.report().unwrap().current_host_bytes, 0);
655    }
656
657    #[test]
658    fn manager_identity_and_accounting_fail_closed() {
659        let pool = pool();
660        let _membership = pool.register_manager(9).unwrap();
661        assert_eq!(
662            pool.register_manager(9).unwrap_err(),
663            CachePoolError::DuplicateManager { manager: 9 }
664        );
665        assert_eq!(
666            pool.update_manager(10, CachePoolUsage::default())
667                .unwrap_err(),
668            CachePoolError::UnknownManager { manager: 10 }
669        );
670        let overflow = CacheResidencyPool::new(CachePoolLimits::new(u64::MAX, 1, 1, 0).unwrap());
671        let _reservation = overflow
672            .reserve(CachePoolUsage {
673                device_bytes: u64::MAX,
674                ..CachePoolUsage::default()
675            })
676            .unwrap();
677        assert!(matches!(
678            overflow.reserve(CachePoolUsage {
679                device_bytes: 1,
680                ..CachePoolUsage::default()
681            }),
682            Err(CachePoolError::AccountingOverflow { .. })
683        ));
684    }
685}
eredu_runtime/cache.rs

eredu_runtime/
cache.rs
