Skip to main content

eredu_runtime/cache/
policy.rs

1//! Backend-neutral mutable-cache residency policy.
2
3use std::{path::PathBuf, sync::Arc};
4
5use eredu_core::residency::CacheEvictionPolicy;
6
7use super::{CachePoolError, CachePoolLimits, CacheResidencyPool};
8
9/// Selects fully resident mutable state or bounded block residency.
10#[derive(Debug, Clone, Default, Eq, PartialEq)]
11pub enum CacheResidencyPolicy {
12    /// Keep state entirely in backend execution memory.
13    #[default]
14    Device,
15    /// Store sealed state in token-addressable blocks under finite budgets.
16    Paged(PagedCacheOptions),
17}
18
19/// Controls optional disk backing for a live inference cache.
20#[derive(Debug, Clone, Default, Eq, PartialEq)]
21pub enum LiveCacheDiskPolicy {
22    /// Do not write live mutable state to disk.
23    #[default]
24    Disabled,
25    /// Retain demoted sealed blocks in an explicit ephemeral directory.
26    Enabled {
27        /// Directory dedicated to this live cache.
28        directory: PathBuf,
29        /// Finite logical byte limit for live cache files.
30        budget_bytes: u64,
31        /// Bound on pending reader or writer requests.
32        queue_capacity: usize,
33    },
34}
35
36/// Validated finite limits for block-addressable mutable state.
37#[derive(Debug, Clone, Eq, PartialEq)]
38pub struct PagedCacheOptions {
39    block_size_tokens: i32,
40    device_budget_bytes: u64,
41    host_budget_bytes: u64,
42    recent_device_blocks: usize,
43    eviction_policy: CacheEvictionPolicy,
44    full_attention: bool,
45    retain_discarded_for_persistence: bool,
46    live_disk: LiveCacheDiskPolicy,
47    sample_process: bool,
48    pool: Option<Arc<CacheResidencyPool>>,
49}
50
51impl PagedCacheOptions {
52    /// Creates paged-state limits. Every memory limit is finite and explicit.
53    pub fn new(
54        block_size_tokens: i32,
55        device_budget_bytes: u64,
56        host_budget_bytes: u64,
57        recent_device_blocks: usize,
58    ) -> Result<Self, CacheResidencyConfigurationError> {
59        if block_size_tokens <= 0 {
60            return Err(CacheResidencyConfigurationError::InvalidOptions(
61                "cache block size must be positive".into(),
62            ));
63        }
64        if device_budget_bytes == 0 {
65            return Err(CacheResidencyConfigurationError::InvalidOptions(
66                "paged cache device budget must be nonzero".into(),
67            ));
68        }
69        if recent_device_blocks == 0 {
70            return Err(CacheResidencyConfigurationError::InvalidOptions(
71                "paged cache must protect at least one recent device block".into(),
72            ));
73        }
74        Ok(Self {
75            block_size_tokens,
76            device_budget_bytes,
77            host_budget_bytes,
78            recent_device_blocks,
79            eviction_policy: CacheEvictionPolicy::LeastRecentlyUsed,
80            full_attention: false,
81            retain_discarded_for_persistence: false,
82            live_disk: LiveCacheDiskPolicy::Disabled,
83            sample_process: false,
84            pool: None,
85        })
86    }
87
88    /// Enables exact blockwise full-context attention.
89    pub const fn with_full_attention(mut self, enabled: bool) -> Self {
90        self.full_attention = enabled;
91        self
92    }
93
94    /// Retains blocks older than a sliding window solely for later persistence.
95    pub const fn with_persistence_retention(mut self, enabled: bool) -> Self {
96        self.retain_discarded_for_persistence = enabled;
97        self
98    }
99
100    /// Selects deterministic block eviction ordering.
101    pub const fn with_eviction_policy(mut self, policy: CacheEvictionPolicy) -> Self {
102        self.eviction_policy = policy;
103        self
104    }
105
106    /// Configures explicit live disk backing.
107    pub fn with_live_disk(
108        mut self,
109        directory: impl Into<PathBuf>,
110        budget_bytes: u64,
111        queue_capacity: usize,
112    ) -> Result<Self, CacheResidencyConfigurationError> {
113        if budget_bytes == 0 {
114            return Err(CacheResidencyConfigurationError::InvalidOptions(
115                "live cache disk budget must be nonzero".into(),
116            ));
117        }
118        if queue_capacity == 0 {
119            return Err(CacheResidencyConfigurationError::InvalidOptions(
120                "live cache disk queue capacity must be nonzero".into(),
121            ));
122        }
123        let directory = directory.into();
124        if directory.as_os_str().is_empty() {
125            return Err(CacheResidencyConfigurationError::InvalidOptions(
126                "live cache disk directory must not be empty".into(),
127            ));
128        }
129        self.live_disk = LiveCacheDiskPolicy::Enabled {
130            directory,
131            budget_bytes,
132            queue_capacity,
133        };
134        Ok(self)
135    }
136
137    /// Enables optional process-memory sampling in reports.
138    pub const fn with_process_sampling(mut self, enabled: bool) -> Self {
139        self.sample_process = enabled;
140        self
141    }
142
143    /// Attaches this per-cache policy to an aggregate process pool.
144    pub fn with_pool(
145        mut self,
146        pool: CacheResidencyPool,
147    ) -> Result<Self, CacheResidencyConfigurationError> {
148        if self.device_budget_bytes > pool.limits().device_bytes() {
149            return Err(CacheResidencyConfigurationError::InvalidOptions(format!(
150                "per-cache device budget {} exceeds cache pool budget {}",
151                self.device_budget_bytes,
152                pool.limits().device_bytes()
153            )));
154        }
155        if self.host_budget_bytes > pool.limits().host_bytes() {
156            return Err(CacheResidencyConfigurationError::InvalidOptions(format!(
157                "per-cache host budget {} exceeds cache pool budget {}",
158                self.host_budget_bytes,
159                pool.limits().host_bytes()
160            )));
161        }
162        if let LiveCacheDiskPolicy::Enabled { budget_bytes, .. } = &self.live_disk {
163            if *budget_bytes > pool.limits().disk_bytes() {
164                return Err(CacheResidencyConfigurationError::InvalidOptions(format!(
165                    "per-cache disk budget {budget_bytes} exceeds cache pool budget {}",
166                    pool.limits().disk_bytes()
167                )));
168            }
169        }
170        self.pool = Some(Arc::new(pool));
171        Ok(self)
172    }
173
174    /// Returns the block size in tokens.
175    pub const fn block_size_tokens(&self) -> i32 {
176        self.block_size_tokens
177    }
178
179    /// Returns the finite logical device-cache budget.
180    pub const fn device_budget_bytes(&self) -> u64 {
181        self.device_budget_bytes
182    }
183
184    /// Returns the finite physical host-transfer allocation budget.
185    pub const fn host_budget_bytes(&self) -> u64 {
186        self.host_budget_bytes
187    }
188
189    /// Returns the recent block count protected on the execution device per layer.
190    pub const fn recent_device_blocks(&self) -> usize {
191        self.recent_device_blocks
192    }
193
194    /// Returns deterministic block eviction ordering.
195    pub const fn eviction_policy(&self) -> CacheEvictionPolicy {
196        self.eviction_policy
197    }
198
199    /// Returns whether exact blockwise full attention is enabled.
200    pub const fn full_attention_enabled(&self) -> bool {
201        self.full_attention
202    }
203
204    /// Returns whether discarded sliding state is retained for persistence.
205    pub const fn retains_discarded_for_persistence(&self) -> bool {
206        self.retain_discarded_for_persistence
207    }
208
209    /// Returns the live disk policy.
210    pub const fn live_disk_policy(&self) -> &LiveCacheDiskPolicy {
211        &self.live_disk
212    }
213
214    /// Returns whether process-memory sampling is enabled.
215    pub const fn process_sampling_enabled(&self) -> bool {
216        self.sample_process
217    }
218
219    /// Returns the explicitly attached aggregate pool, if any.
220    pub fn pool(&self) -> Option<&CacheResidencyPool> {
221        self.pool.as_deref()
222    }
223
224    /// Creates the default aggregate ownership pool for this policy.
225    pub fn create_pool(&self) -> Result<CacheResidencyPool, CacheResidencyConfigurationError> {
226        let disk_bytes = match self.live_disk_policy() {
227            LiveCacheDiskPolicy::Disabled => 0,
228            LiveCacheDiskPolicy::Enabled { budget_bytes, .. } => *budget_bytes,
229        };
230        let transfer_bytes = self
231            .device_budget_bytes()
232            .max(self.host_budget_bytes())
233            .saturating_mul(2)
234            .max(1);
235        Ok(CacheResidencyPool::new(CachePoolLimits::new(
236            self.device_budget_bytes(),
237            self.host_budget_bytes(),
238            transfer_bytes,
239            disk_bytes,
240        )?))
241    }
242}
243
244/// Invalid backend-neutral mutable-cache residency configuration.
245#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
246pub enum CacheResidencyConfigurationError {
247    /// Paged options were contradictory or unbounded.
248    #[error("invalid paged cache options: {0}")]
249    InvalidOptions(String),
250    /// Aggregate pool limits were invalid.
251    #[error(transparent)]
252    Pool(#[from] CachePoolError),
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn finite_limits_and_live_disk_are_validated_without_a_backend() {
261        assert!(PagedCacheOptions::new(0, 1, 1, 1).is_err());
262        assert!(PagedCacheOptions::new(16, 0, 1, 1).is_err());
263        assert!(PagedCacheOptions::new(16, 1, 1, 0).is_err());
264        let options = PagedCacheOptions::new(16, 8, 0, 1)
265            .unwrap()
266            .with_live_disk("cache", 32, 2)
267            .unwrap();
268        assert_eq!(options.create_pool().unwrap().limits().disk_bytes(), 32);
269    }
270
271    #[test]
272    fn attached_pool_must_cover_every_enabled_residency_tier() {
273        let pool = CacheResidencyPool::new(CachePoolLimits::new(8, 4, 8, 0).unwrap());
274        let error = PagedCacheOptions::new(16, 16, 4, 1)
275            .unwrap()
276            .with_pool(pool)
277            .unwrap_err();
278        assert!(matches!(
279            error,
280            CacheResidencyConfigurationError::InvalidOptions(_)
281        ));
282    }
283}