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

eredu_runtime/
cache.rs
