Skip to main content

car_inference/
resource_policy.rs

1//! Persisted admission policy for CAR-managed local model memory.
2//!
3//! This module answers how much host memory CAR may admit for local models. It
4//! deliberately does not claim to impose an operating-system RSS limit. Apple
5//! Silicon and CPU inference use the host budget; discrete CUDA keeps measured
6//! VRAM as a separate fit constraint instead of applying the Apple unified-
7//! memory percentage to it.
8
9use std::collections::{HashMap, HashSet};
10use std::io::{self, Read, Write};
11use std::path::{Path, PathBuf};
12use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
13use std::sync::{Arc, Weak};
14use std::sync::{Mutex, OnceLock};
15
16use serde::{Deserialize, Serialize};
17use thiserror::Error;
18
19use crate::hardware::{GpuBackend, HardwareInfo};
20use crate::schema::ModelSchema;
21
22/// File name under the CAR state root.
23pub const RESOURCE_POLICY_FILE: &str = "model-resource-policy.json";
24
25const EVERYDAY_MODEL_PERCENT: u64 = 40;
26const LOCAL_FOCUSED_MODEL_PERCENT: u64 = 80;
27const EMERGENCY_RESERVE_PERCENT: u64 = 10;
28const MINIMUM_EMERGENCY_RESERVE_MB: u64 = 2 * 1024;
29const MAX_POLICY_BYTES: u64 = 64 * 1024;
30
31/// Context window used by recommendation-time memory estimates. This is a
32/// typical working window, not a model's advertised maximum.
33pub const RECOMMENDATION_CONTEXT_TOKENS: usize = 8_192;
34
35const METAL_RUNTIME_OVERHEAD_MB: u64 = 512;
36const CUDA_RUNTIME_OVERHEAD_MB: u64 = 512;
37const CPU_RUNTIME_OVERHEAD_MB: u64 = 1_024;
38const TRANSIENT_ALLOCATION_MARGIN_MB: u64 = 1_024;
39/// Memory accounting uses binary megabytes throughout this module.
40const BYTES_PER_MB: u128 = 1024 * 1024;
41
42/// Preset or exact user-selected local-model admission behavior.
43#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
44#[serde(rename_all = "snake_case")]
45pub enum ResourceProfile {
46    Everyday,
47    LocalFocused,
48    Custom,
49}
50
51/// The persisted local-model memory preference.
52#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
53#[serde(deny_unknown_fields)]
54pub struct ResourcePolicy {
55    pub profile: ResourceProfile,
56    pub custom_max_model_mb: Option<u64>,
57}
58
59impl Default for ResourcePolicy {
60    fn default() -> Self {
61        Self::everyday()
62    }
63}
64
65impl ResourcePolicy {
66    pub fn everyday() -> Self {
67        Self {
68            profile: ResourceProfile::Everyday,
69            custom_max_model_mb: None,
70        }
71    }
72
73    pub fn local_focused() -> Self {
74        Self {
75            profile: ResourceProfile::LocalFocused,
76            custom_max_model_mb: None,
77        }
78    }
79
80    /// Build an exact Custom policy from user-facing GB input.
81    ///
82    /// Only finite, nonnegative half-GB increments are accepted. The persisted
83    /// value is integer MB, so it round-trips without floating-point drift.
84    pub fn custom_gb(gigabytes: f64) -> Result<Self, ResourcePolicyError> {
85        if !gigabytes.is_finite() || gigabytes < 0.0 {
86            return Err(ResourcePolicyError::InvalidCustomGigabytes(gigabytes));
87        }
88
89        let half_gb_steps = gigabytes * 2.0;
90        if half_gb_steps.fract() != 0.0 || half_gb_steps > u64::MAX as f64 {
91            return Err(ResourcePolicyError::InvalidCustomGigabytes(gigabytes));
92        }
93        let steps = half_gb_steps as u64;
94        let custom_max_model_mb = steps
95            .checked_mul(512)
96            .ok_or(ResourcePolicyError::InvalidCustomGigabytes(gigabytes))?;
97
98        Ok(Self {
99            profile: ResourceProfile::Custom,
100            custom_max_model_mb: Some(custom_max_model_mb),
101        })
102    }
103
104    /// Evaluate the stored preference against the current machine capacity.
105    ///
106    /// Custom values are clamped for this evaluation only. The persisted value
107    /// is left untouched so a preference moved between machines is never
108    /// silently rewritten.
109    pub fn effective_budget(&self, total_memory_mb: u64) -> EffectiveResourceBudget {
110        let emergency_reserve_mb = minimum_emergency_reserve(total_memory_mb);
111        let safe_maximum_mb = total_memory_mb.saturating_sub(emergency_reserve_mb);
112        let requested_ceiling_mb = match self.profile {
113            ResourceProfile::Everyday => percent_of(total_memory_mb, EVERYDAY_MODEL_PERCENT),
114            ResourceProfile::LocalFocused => {
115                percent_of(total_memory_mb, LOCAL_FOCUSED_MODEL_PERCENT)
116            }
117            ResourceProfile::Custom => self.custom_max_model_mb.unwrap_or(0),
118        };
119        let configured_model_ceiling_mb = requested_ceiling_mb.min(safe_maximum_mb);
120        let normalization_notice = (matches!(self.profile, ResourceProfile::Custom)
121            && requested_ceiling_mb > safe_maximum_mb)
122            .then(|| {
123                format!(
124                    "The saved Custom allocation was adjusted from {requested_ceiling_mb} MB to \
125                     {safe_maximum_mb} MB on this machine to preserve the \
126                     {emergency_reserve_mb} MB emergency reserve."
127                )
128            });
129
130        EffectiveResourceBudget {
131            total_memory_mb,
132            emergency_reserve_mb,
133            configured_model_ceiling_mb,
134            effective_new_load_ceiling_mb: configured_model_ceiling_mb,
135            normalization_notice,
136        }
137    }
138
139    /// Recommendation target within the operating ceiling.
140    ///
141    /// Everyday intentionally targets half of its 40% admission ceiling so
142    /// the primary leaves room for context growth, CAR, and changing app
143    /// load. Other profiles may use their full configured ceiling.
144    pub fn recommendation_target_mb(&self, total_memory_mb: u64) -> u64 {
145        let ceiling = self
146            .effective_budget(total_memory_mb)
147            .configured_model_ceiling_mb;
148        if self.profile == ResourceProfile::Everyday {
149            ceiling / 2
150        } else {
151            ceiling
152        }
153    }
154
155    /// Validate the exact persisted shape accepted by CAR.
156    ///
157    /// Keeping this check on the value itself ensures file reads and writes
158    /// cannot drift into accepting different policy documents.
159    pub fn validate(&self) -> Result<(), ResourcePolicyError> {
160        match self.profile {
161            ResourceProfile::Custom => match self.custom_max_model_mb {
162                Some(value) if value.is_multiple_of(512) => Ok(()),
163                Some(value) => Err(ResourcePolicyError::InvalidPolicy {
164                    reason: format!(
165                        "Custom model RAM must be a 0.5 GB (512 MB) increment; got {value} MB"
166                    ),
167                }),
168                None => Err(ResourcePolicyError::InvalidPolicy {
169                    reason: "Custom profile requires custom_max_model_mb".into(),
170                }),
171            },
172            ResourceProfile::Everyday | ResourceProfile::LocalFocused => {
173                if self.custom_max_model_mb.is_none() {
174                    Ok(())
175                } else {
176                    Err(ResourcePolicyError::InvalidPolicy {
177                        reason: format!(
178                            "{:?} profile must not set custom_max_model_mb",
179                            self.profile
180                        ),
181                    })
182                }
183            }
184        }
185    }
186}
187
188/// Static policy math. Live-memory evidence is added by the admission service.
189#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
190pub struct EffectiveResourceBudget {
191    pub total_memory_mb: u64,
192    pub emergency_reserve_mb: u64,
193    pub configured_model_ceiling_mb: u64,
194    pub effective_new_load_ceiling_mb: u64,
195    pub normalization_notice: Option<String>,
196}
197
198/// Measured dedicated accelerator memory. This is intentionally separate from
199/// host RAM; CAR does not apply Apple unified-memory percentages to CUDA VRAM.
200#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
201pub struct AcceleratorResourceBudget {
202    pub total_mb: u64,
203    pub budget_mb: u64,
204}
205
206/// Host policy plus an optional distinct accelerator fit pool.
207#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
208pub struct ResourceEvaluation {
209    pub host_memory: EffectiveResourceBudget,
210    pub accelerator_memory: Option<AcceleratorResourceBudget>,
211}
212
213/// Why CAR believes the weight portion of a model-memory estimate.
214#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
215#[serde(rename_all = "snake_case")]
216pub enum ModelResourceEvidence {
217    /// The catalog declared a RAM or artifact-size floor.
218    CatalogExact,
219    /// A caller supplied a physical on-disk measurement.
220    FileSystemMeasured,
221    /// CAR estimated from parameter count because no stronger fact existed.
222    Heuristic,
223}
224
225/// Recommendation-time cold-load estimate. This is static capacity evidence,
226/// not a live-memory probe or an operating-system RSS guarantee.
227#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
228pub struct ModelMemoryEstimate {
229    pub weights_mb: u64,
230    pub runtime_overhead_mb: u64,
231    pub context_overhead_mb: u64,
232    pub transient_margin_mb: u64,
233    pub estimated_peak_mb: u64,
234    pub evidence: ModelResourceEvidence,
235}
236
237/// Result of checking one CAR-managed local allocation against both the saved
238/// policy and live machine headroom. `LiveMemoryUnknown` is intentionally not
239/// collapsed into `Allowed`: callers may choose a conservative static-policy
240/// fallback, but must not present the model as safe *right now*.
241#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
242pub struct LocalLoadPreflight {
243    pub model_id: String,
244    pub estimate: ModelMemoryEstimate,
245    pub configured_ceiling_mb: u64,
246    pub resident_model_mb: u64,
247    pub active_reservations_mb: u64,
248    pub estimated_incremental_mb: u64,
249    pub accelerator_total_mb: Option<u64>,
250    pub accelerator_resident_mb: Option<u64>,
251    pub accelerator_incremental_mb: Option<u64>,
252    pub live_available_mb: Option<u64>,
253    pub emergency_reserve_mb: u64,
254    pub verdict: LocalLoadVerdict,
255}
256
257#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
258#[serde(rename_all = "snake_case")]
259pub enum LocalLoadVerdict {
260    Allowed,
261    LiveMemoryUnknown,
262    DisabledByPolicy,
263    ExceedsConfiguredCeiling,
264    InsufficientLiveMemory,
265    ModelMaintenance,
266    /// A CAR-owned process is still being reaped. Its weights remain charged
267    /// and no replacement may peer-discount them until exit is confirmed.
268    PendingTeardown,
269}
270
271impl LocalLoadVerdict {
272    pub fn permits_static_fallback(&self) -> bool {
273        matches!(self, Self::Allowed | Self::LiveMemoryUnknown)
274    }
275}
276
277/// Machine-specific live-memory evidence. Returning `Ok(None)` is distinct
278/// from zero available memory and is preserved in the public preflight.
279pub trait LiveMemoryProbe: Send + Sync {
280    fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError>;
281}
282
283#[derive(Default)]
284pub struct SystemLiveMemoryProbe;
285
286impl LiveMemoryProbe for SystemLiveMemoryProbe {
287    fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
288        Ok(crate::hardware::available_ram_mb())
289    }
290}
291
292#[cfg(test)]
293pub(crate) struct FixedLiveMemoryProbe(Option<u64>);
294
295#[cfg(test)]
296impl FixedLiveMemoryProbe {
297    pub(crate) fn known(available_mb: u64) -> Self {
298        Self(Some(available_mb))
299    }
300
301    fn unknown() -> Self {
302        Self(None)
303    }
304}
305
306#[cfg(test)]
307impl LiveMemoryProbe for FixedLiveMemoryProbe {
308    fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
309        Ok(self.0)
310    }
311}
312
313#[derive(Clone, Copy, Debug, PartialEq, Eq)]
314enum WeightPlacement {
315    Host,
316    Accelerator,
317}
318
319#[derive(Clone, Debug)]
320struct ResidentAllocation {
321    weights_mb: u64,
322    placement: WeightPlacement,
323    logical_model_id: String,
324}
325
326#[derive(Clone, Default)]
327struct AdmissionState {
328    resident_models: HashMap<String, ResidentAllocation>,
329    active_host_reservations_mb: u64,
330    active_accelerator_reservations_mb: u64,
331    active_by_model: HashMap<String, usize>,
332    maintenance_models: HashSet<String>,
333    pending_teardown_models: HashMap<String, HashSet<String>>,
334    model_aliases: HashMap<String, HashSet<String>>,
335    next_request_id: u64,
336}
337
338#[derive(Clone, Default)]
339struct MachineAdmissionLedger {
340    resident_models: HashMap<(u64, String), ResidentAllocation>,
341    pending_allocations: HashMap<(u64, String), ResidentAllocation>,
342    active_host_by_owner: HashMap<u64, u64>,
343    active_accelerator_by_owner: HashMap<u64, u64>,
344}
345
346fn process_machine_admission_ledger() -> Arc<Mutex<MachineAdmissionLedger>> {
347    static LEDGER: OnceLock<Arc<Mutex<MachineAdmissionLedger>>> = OnceLock::new();
348    LEDGER
349        .get_or_init(|| Arc::new(Mutex::new(MachineAdmissionLedger::default())))
350        .clone()
351}
352
353fn next_admission_owner_id() -> u64 {
354    static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
355    NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
356}
357
358/// Allocation identities are source-qualified so a worker process and a
359/// supervised server for the same logical model cannot overwrite or evict one
360/// another's physical-memory record.
361pub fn worker_process_allocation_id(model_id: &str) -> String {
362    format!("worker:{model_id}")
363}
364
365pub fn vllm_process_allocation_id(model_id: &str) -> String {
366    format!("vllm:{model_id}")
367}
368
369/// One process-wide coordinator for local cold loads and per-request runtime
370/// allocations. Its mutex makes check-and-reserve one atomic operation, so two
371/// individually-safe requests cannot both consume the same headroom.
372pub struct LocalAdmissionCoordinator {
373    policy: std::sync::RwLock<ResourcePolicy>,
374    hardware: HardwareInfo,
375    live_probe: Arc<dyn LiveMemoryProbe>,
376    state: Mutex<AdmissionState>,
377    machine_ledger: Arc<Mutex<MachineAdmissionLedger>>,
378    /// Cross-process activity ownership for exact resident allocation ids.
379    /// Kept outside cloneable accounting snapshots so an idle cache/process
380    /// holds its shared lease until exact eviction or confirmed teardown.
381    resident_activity_leases: Mutex<HashMap<String, Arc<crate::model_management::ModelLease>>>,
382    models_dir: std::sync::RwLock<Option<PathBuf>>,
383    kv_geometry_cache: Mutex<HashMap<String, (ConfigStamp, KvGeometry)>>,
384    owner_id: u64,
385}
386
387impl Drop for LocalAdmissionCoordinator {
388    fn drop(&mut self) {
389        let mut machine = self
390            .machine_ledger
391            .lock()
392            .unwrap_or_else(std::sync::PoisonError::into_inner);
393        machine
394            .resident_models
395            .retain(|(owner_id, _), _| *owner_id != self.owner_id);
396        machine
397            .pending_allocations
398            .retain(|(owner_id, _), _| *owner_id != self.owner_id);
399        machine.active_host_by_owner.remove(&self.owner_id);
400        machine.active_accelerator_by_owner.remove(&self.owner_id);
401    }
402}
403
404fn scoped_admission_registry() -> &'static Mutex<HashMap<PathBuf, Weak<LocalAdmissionCoordinator>>>
405{
406    static REGISTRY: OnceLock<Mutex<HashMap<PathBuf, Weak<LocalAdmissionCoordinator>>>> =
407        OnceLock::new();
408    REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
409}
410
411/// Stable identity for a state root. Existing roots are filesystem-
412/// canonicalized so symlink/relative aliases cannot split admission accounting.
413/// If canonicalization is unavailable (commonly before first-run directory
414/// creation), lexical absolute normalization is deliberately conservative:
415/// aliases may be charged independently, but never peer-discount one another.
416pub fn normalized_state_root_key(state_root: &Path) -> PathBuf {
417    let absolute = if state_root.is_absolute() {
418        state_root.to_path_buf()
419    } else if let Ok(current) = std::env::current_dir() {
420        current.join(state_root)
421    } else {
422        return state_root.to_path_buf();
423    };
424    // Do not lexically collapse `..` before asking the filesystem. In
425    // `symlink/../state`, `..` applies after resolving the symlink and can name
426    // a different physical directory than a lexical pop would produce.
427    if let Ok(canonical) = std::fs::canonicalize(&absolute) {
428        return canonical;
429    }
430
431    // First-run roots often do not exist yet. Canonicalize the nearest
432    // existing ancestor so a missing leaf beneath a symlink still shares the
433    // same coordinator/runtime identity as its real path.
434    let mut ancestor = absolute.clone();
435    let mut missing_suffix = Vec::new();
436    while let Some(component) = ancestor.components().next_back() {
437        let name = match component {
438            std::path::Component::Normal(name) => name.to_owned(),
439            std::path::Component::CurDir => std::ffi::OsString::from("."),
440            std::path::Component::ParentDir => std::ffi::OsString::from(".."),
441            std::path::Component::RootDir | std::path::Component::Prefix(_) => break,
442        };
443        missing_suffix.push(name);
444        ancestor.pop();
445        if let Ok(mut canonical) = std::fs::canonicalize(&ancestor) {
446            for component in missing_suffix.iter().rev() {
447                if component == std::ffi::OsStr::new(".") {
448                    continue;
449                }
450                if component == std::ffi::OsStr::new("..") {
451                    canonical.pop();
452                } else {
453                    canonical.push(component);
454                }
455            }
456            return canonical;
457        }
458    }
459    // No ancestor could be canonicalized. Keep the absolute spelling rather
460    // than inventing peer identity; conservative duplicate charging is safer.
461    absolute
462}
463
464pub fn install_shared_local_admission(coordinator: Arc<LocalAdmissionCoordinator>) {
465    let root = normalized_state_root_key(&car_home::root_or_relative());
466    scoped_admission_registry()
467        .lock()
468        .unwrap_or_else(std::sync::PoisonError::into_inner)
469        .entry(root)
470        .or_insert_with(|| Arc::downgrade(&coordinator));
471}
472
473pub fn shared_local_admission() -> Arc<LocalAdmissionCoordinator> {
474    let root = car_home::root_or_relative();
475    let policy = FileResourcePolicyRepository::new(root.clone())
476        .load()
477        .unwrap_or_else(|_| ResourcePolicy::everyday());
478    scoped_local_admission(root, policy, HardwareInfo::detect())
479}
480
481pub fn scoped_local_admission(
482    state_root: impl AsRef<Path>,
483    policy: ResourcePolicy,
484    hardware: HardwareInfo,
485) -> Arc<LocalAdmissionCoordinator> {
486    scoped_local_admission_with_models_dir(state_root, policy, hardware, None)
487}
488
489/// Return the admission coordinator for a state root, with optional access to
490/// installed model configs for architecture-aware context-memory estimates.
491pub fn scoped_local_admission_with_models_dir(
492    state_root: impl AsRef<Path>,
493    policy: ResourcePolicy,
494    hardware: HardwareInfo,
495    models_dir: Option<PathBuf>,
496) -> Arc<LocalAdmissionCoordinator> {
497    let state_root = normalized_state_root_key(state_root.as_ref());
498    let mut registry = scoped_admission_registry()
499        .lock()
500        .unwrap_or_else(std::sync::PoisonError::into_inner);
501    registry.retain(|_, coordinator| coordinator.strong_count() > 0);
502    if let Some(existing) = registry.get(&state_root).and_then(Weak::upgrade) {
503        existing.set_policy(policy);
504        if let Some(models_dir) = models_dir {
505            existing.set_models_dir(models_dir);
506        }
507        return existing;
508    }
509    let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe_ledger_and_models_dir(
510        policy,
511        hardware,
512        Arc::new(SystemLiveMemoryProbe),
513        process_machine_admission_ledger(),
514        models_dir,
515    ));
516    registry.insert(state_root, Arc::downgrade(&coordinator));
517    coordinator
518}
519
520pub fn local_admission_for_scope(
521    state_root: impl AsRef<Path>,
522) -> Option<Arc<LocalAdmissionCoordinator>> {
523    let state_root = normalized_state_root_key(state_root.as_ref());
524    scoped_admission_registry()
525        .lock()
526        .unwrap_or_else(std::sync::PoisonError::into_inner)
527        .get(&state_root)
528        .and_then(Weak::upgrade)
529}
530
531impl LocalAdmissionCoordinator {
532    fn built_in_model_identity(model_id: &str) -> Option<String> {
533        let lower = model_id.to_ascii_lowercase();
534        if lower.contains("kokoro-82m") {
535            if lower.contains("6bit") {
536                return Some("mlx/kokoro-82m:6bit".into());
537            }
538            if lower.contains("bf16") {
539                return Some("mlx/kokoro-82m:bf16".into());
540            }
541        }
542        None
543    }
544
545    fn resolve_model_identity(state: &AdmissionState, model_id: &str) -> String {
546        state
547            .model_aliases
548            .iter()
549            .find_map(|(canonical, aliases)| {
550                (canonical == model_id || aliases.contains(model_id)).then(|| canonical.clone())
551            })
552            .or_else(|| Self::built_in_model_identity(model_id))
553            .unwrap_or_else(|| model_id.to_string())
554    }
555
556    /// Resolve any registered catalog/artifact/provider alias to the single
557    /// allocation key used by admission, caches, and maintenance.
558    pub fn canonical_model_id(&self, model_id: &str) -> String {
559        let state = self
560            .state
561            .lock()
562            .unwrap_or_else(std::sync::PoisonError::into_inner);
563        Self::resolve_model_identity(&state, model_id)
564    }
565
566    pub fn register_model_aliases<I, S>(&self, canonical_model_id: &str, aliases: I)
567    where
568        I: IntoIterator<Item = S>,
569        S: Into<String>,
570    {
571        let mut aliases = aliases
572            .into_iter()
573            .map(Into::into)
574            .collect::<HashSet<String>>();
575        aliases.insert(canonical_model_id.to_string());
576        let mut state = self
577            .state
578            .lock()
579            .unwrap_or_else(std::sync::PoisonError::into_inner);
580        let mut canonical_model_id = Self::built_in_model_identity(canonical_model_id)
581            .unwrap_or_else(|| canonical_model_id.to_string());
582        let intersecting = state
583            .model_aliases
584            .iter()
585            .filter(|(existing, existing_aliases)| {
586                aliases.contains(*existing)
587                    || existing_aliases.iter().any(|alias| aliases.contains(alias))
588            })
589            .map(|(existing, _)| existing.clone())
590            .collect::<Vec<_>>();
591        if let Some(existing) = intersecting.first() {
592            canonical_model_id = existing.clone();
593        }
594        for existing in intersecting {
595            if let Some(existing_aliases) = state.model_aliases.remove(&existing) {
596                aliases.extend(existing_aliases);
597            }
598            aliases.insert(existing);
599        }
600        aliases.insert(canonical_model_id.clone());
601
602        let mut active = 0usize;
603        for alias in &aliases {
604            active = active.saturating_add(state.active_by_model.remove(alias).unwrap_or_default());
605        }
606        if active > 0 {
607            *state
608                .active_by_model
609                .entry(canonical_model_id.clone())
610                .or_default() += active;
611        }
612        let mut maintenance = false;
613        for alias in &aliases {
614            maintenance |= state.maintenance_models.remove(alias);
615        }
616        if maintenance {
617            state.maintenance_models.insert(canonical_model_id.clone());
618        }
619        let mut pending = HashSet::new();
620        for alias in &aliases {
621            pending.extend(
622                state
623                    .pending_teardown_models
624                    .remove(alias)
625                    .unwrap_or_default(),
626            );
627        }
628        if !pending.is_empty() {
629            state
630                .pending_teardown_models
631                .entry(canonical_model_id.clone())
632                .or_default()
633                .extend(pending);
634        }
635        for resident in state.resident_models.values_mut() {
636            if aliases.contains(&resident.logical_model_id) {
637                resident.logical_model_id = canonical_model_id.clone();
638            }
639        }
640        state
641            .model_aliases
642            .insert(canonical_model_id.clone(), aliases.clone());
643        // Lock order is always scoped state then machine ledger. Holding both
644        // here makes late alias registration one atomic identity migration;
645        // reserve cannot observe half-rekeyed maintenance/residency state.
646        let mut machine = self
647            .machine_ledger
648            .lock()
649            .unwrap_or_else(std::sync::PoisonError::into_inner);
650        for ((owner_id, _), resident) in machine.resident_models.iter_mut() {
651            if *owner_id == self.owner_id && aliases.contains(&resident.logical_model_id) {
652                resident.logical_model_id = canonical_model_id.clone();
653            }
654        }
655        for ((owner_id, _), resident) in machine.pending_allocations.iter_mut() {
656            if *owner_id == self.owner_id && aliases.contains(&resident.logical_model_id) {
657                resident.logical_model_id = canonical_model_id.clone();
658            }
659        }
660    }
661    pub fn new(policy: ResourcePolicy, hardware: HardwareInfo) -> Self {
662        Self::with_probe(policy, hardware, Arc::new(SystemLiveMemoryProbe))
663    }
664
665    pub fn with_probe(
666        policy: ResourcePolicy,
667        hardware: HardwareInfo,
668        live_probe: Arc<dyn LiveMemoryProbe>,
669    ) -> Self {
670        Self::with_probe_and_ledger(
671            policy,
672            hardware,
673            live_probe,
674            Arc::new(Mutex::new(MachineAdmissionLedger::default())),
675        )
676    }
677
678    fn with_probe_and_ledger(
679        policy: ResourcePolicy,
680        hardware: HardwareInfo,
681        live_probe: Arc<dyn LiveMemoryProbe>,
682        machine_ledger: Arc<Mutex<MachineAdmissionLedger>>,
683    ) -> Self {
684        Self::with_probe_ledger_and_models_dir(policy, hardware, live_probe, machine_ledger, None)
685    }
686
687    fn with_probe_ledger_and_models_dir(
688        policy: ResourcePolicy,
689        hardware: HardwareInfo,
690        live_probe: Arc<dyn LiveMemoryProbe>,
691        machine_ledger: Arc<Mutex<MachineAdmissionLedger>>,
692        models_dir: Option<PathBuf>,
693    ) -> Self {
694        Self {
695            policy: std::sync::RwLock::new(policy),
696            hardware,
697            live_probe,
698            state: Mutex::new(AdmissionState::default()),
699            machine_ledger,
700            resident_activity_leases: Mutex::new(HashMap::new()),
701            models_dir: std::sync::RwLock::new(models_dir),
702            kv_geometry_cache: Mutex::new(HashMap::new()),
703            owner_id: next_admission_owner_id(),
704        }
705    }
706
707    #[cfg(test)]
708    fn with_probe_and_models_dir(
709        policy: ResourcePolicy,
710        hardware: HardwareInfo,
711        live_probe: Arc<dyn LiveMemoryProbe>,
712        models_dir: Option<PathBuf>,
713    ) -> Self {
714        Self::with_probe_ledger_and_models_dir(
715            policy,
716            hardware,
717            live_probe,
718            Arc::new(Mutex::new(MachineAdmissionLedger::default())),
719            models_dir,
720        )
721    }
722
723    fn model_memory_estimate(
724        &self,
725        model: &ModelSchema,
726        context_tokens: usize,
727    ) -> ModelMemoryEstimate {
728        let geometry = self.kv_geometry(model);
729        estimate_model_memory_with_geometry(
730            model,
731            &self.hardware,
732            context_tokens,
733            None,
734            geometry.as_ref(),
735        )
736    }
737
738    /// The installed checkpoint's KV geometry under this coordinator's models
739    /// directory, cached per model id together with the `config.json` stamp it
740    /// was read from. Every lookup stats the file: a removed checkpoint drops
741    /// its entry and falls back, and a replaced or re-downloaded one (new
742    /// mtime or length) is read again, so neither the catalog nor admission
743    /// sizes a row from a config that is no longer there.
744    fn kv_geometry(&self, model: &ModelSchema) -> Option<KvGeometry> {
745        let models_dir = self
746            .models_dir
747            .read()
748            .unwrap_or_else(std::sync::PoisonError::into_inner)
749            .clone();
750        let mut cache = self
751            .kv_geometry_cache
752            .lock()
753            .unwrap_or_else(std::sync::PoisonError::into_inner);
754        // Stat before reading: a write landing between the two leaves an older
755        // stamp beside newer contents, which the next lookup reads again.
756        let Some(stamp) = models_dir
757            .as_deref()
758            .and_then(|dir| ConfigStamp::of(&kv_config_path(model, dir)))
759        else {
760            cache.remove(&model.id);
761            return None;
762        };
763        if let Some((cached, geometry)) = cache.get(&model.id) {
764            if *cached == stamp {
765                return Some(geometry.clone());
766            }
767        }
768        let geometry = kv_geometry_for(model, models_dir.as_deref());
769        match geometry.as_ref() {
770            Some(geometry) => {
771                cache.insert(model.id.clone(), (stamp, geometry.clone()));
772            }
773            None => {
774                cache.remove(&model.id);
775            }
776        }
777        geometry
778    }
779
780    /// The geometry admission sizes these rows' KV cache from, for the
781    /// recommendation estimate to use the same input (car#2213).
782    pub fn checkpoint_geometry<'a>(
783        &self,
784        models: impl IntoIterator<Item = &'a ModelSchema>,
785    ) -> CheckpointGeometry {
786        CheckpointGeometry {
787            by_id: models
788                .into_iter()
789                .filter_map(|model| Some((model.id.clone(), self.kv_geometry(model)?)))
790                .collect(),
791        }
792    }
793
794    fn set_models_dir(&self, models_dir: PathBuf) {
795        let mut configured = self
796            .models_dir
797            .write()
798            .unwrap_or_else(std::sync::PoisonError::into_inner);
799        if configured.as_ref() == Some(&models_dir) {
800            return;
801        }
802        *configured = Some(models_dir);
803        self.kv_geometry_cache
804            .lock()
805            .unwrap_or_else(std::sync::PoisonError::into_inner)
806            .clear();
807    }
808
809    pub fn set_policy(&self, policy: ResourcePolicy) {
810        *self
811            .policy
812            .write()
813            .unwrap_or_else(std::sync::PoisonError::into_inner) = policy;
814    }
815
816    pub fn policy(&self) -> ResourcePolicy {
817        self.policy
818            .read()
819            .unwrap_or_else(std::sync::PoisonError::into_inner)
820            .clone()
821    }
822
823    pub fn mark_resident(&self, model_id: &str, weights_mb: u64) {
824        let placement = self.default_weight_placement();
825        self.mark_resident_with_placement(model_id, model_id, weights_mb, placement);
826    }
827
828    pub fn mark_resident_allocation(
829        &self,
830        logical_model_id: &str,
831        allocation_id: &str,
832        weights_mb: u64,
833    ) {
834        let placement = self.default_weight_placement();
835        self.mark_resident_with_placement(logical_model_id, allocation_id, weights_mb, placement);
836    }
837
838    pub fn mark_evicted(&self, model_id: &str) {
839        let mut state = self
840            .state
841            .lock()
842            .unwrap_or_else(std::sync::PoisonError::into_inner);
843        // Eviction is allocation-owner exact. Logical-model-wide removal can
844        // erase a sibling cache/process/voice lease that still owns weights.
845        state.resident_models.remove(model_id);
846        self.machine_ledger
847            .lock()
848            .unwrap_or_else(std::sync::PoisonError::into_inner)
849            .resident_models
850            .remove(&(self.owner_id, model_id.to_string()));
851        self.resident_activity_leases
852            .lock()
853            .unwrap_or_else(std::sync::PoisonError::into_inner)
854            .remove(model_id);
855    }
856
857    /// Enter fail-closed process teardown. Residency remains charged while a
858    /// background owner retains and reaps the child.
859    pub fn mark_teardown_pending(&self, model_id: &str) {
860        self.mark_teardown_pending_allocation(model_id, model_id);
861    }
862
863    /// Quarantine one concrete process allocation while retaining the logical
864    /// model identity used by maintenance and admission. Multiple process
865    /// owners for one model are reference-counted independently.
866    pub fn mark_teardown_pending_allocation(&self, logical_model_id: &str, allocation_id: &str) {
867        self.mark_teardown_pending_allocation_with_charge(logical_model_id, allocation_id, 0);
868    }
869
870    /// Transfer an exact process allocation into quarantine. If it has not yet
871    /// published residency (the pre-ACK cancellation case), `measured_bytes`
872    /// supplies the physical charge that must outlive the caller reservation.
873    /// Repeating this for the same allocation is idempotent.
874    pub fn mark_teardown_pending_allocation_with_charge(
875        &self,
876        logical_model_id: &str,
877        allocation_id: &str,
878        measured_bytes: u64,
879    ) {
880        let mut state = self
881            .state
882            .lock()
883            .unwrap_or_else(std::sync::PoisonError::into_inner);
884        let model_id = Self::resolve_model_identity(&state, logical_model_id);
885        state
886            .pending_teardown_models
887            .entry(model_id.clone())
888            .or_default()
889            .insert(allocation_id.to_string());
890        let mut machine = self
891            .machine_ledger
892            .lock()
893            .unwrap_or_else(std::sync::PoisonError::into_inner);
894        let key = (self.owner_id, allocation_id.to_string());
895        let machine_resident = machine.resident_models.remove(&key);
896        let scoped_resident = state.resident_models.get(allocation_id).cloned();
897        let resident = machine_resident.or(scoped_resident);
898        let measured_mb = measured_bytes.div_ceil(1024 * 1024);
899        let placement = self.default_weight_placement();
900        machine
901            .pending_allocations
902            .entry(key)
903            .and_modify(|pending| {
904                pending.weights_mb = pending.weights_mb.max(measured_mb);
905                pending.logical_model_id = model_id.clone();
906            })
907            .or_insert_with(|| {
908                resident.unwrap_or(ResidentAllocation {
909                    weights_mb: measured_mb,
910                    placement,
911                    logical_model_id: model_id,
912                })
913            });
914    }
915
916    /// Clear a quarantined allocation only after its OS process was reaped.
917    pub fn finish_teardown(&self, model_id: &str) {
918        self.finish_teardown_allocation(model_id, model_id);
919    }
920
921    /// Finish one concrete process allocation after confirmed OS exit. This
922    /// removes only that owner's resident record; sibling allocations for the
923    /// same logical model remain charged and pending until their own ACK.
924    pub fn finish_teardown_allocation(&self, logical_model_id: &str, allocation_id: &str) {
925        let mut state = self
926            .state
927            .lock()
928            .unwrap_or_else(std::sync::PoisonError::into_inner);
929        let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
930        let exact_pending = state
931            .pending_teardown_models
932            .get_mut(&logical_model_id)
933            .is_some_and(|pending| pending.remove(allocation_id));
934        if exact_pending
935            && state
936                .pending_teardown_models
937                .get(&logical_model_id)
938                .is_some_and(HashSet::is_empty)
939        {
940            state.pending_teardown_models.remove(&logical_model_id);
941        }
942        if !exact_pending {
943            return;
944        }
945        state.resident_models.remove(allocation_id);
946        let mut machine = self
947            .machine_ledger
948            .lock()
949            .unwrap_or_else(std::sync::PoisonError::into_inner);
950        machine
951            .resident_models
952            .remove(&(self.owner_id, allocation_id.to_string()));
953        machine
954            .pending_allocations
955            .remove(&(self.owner_id, allocation_id.to_string()));
956        self.resident_activity_leases
957            .lock()
958            .unwrap_or_else(std::sync::PoisonError::into_inner)
959            .remove(allocation_id);
960    }
961
962    pub fn teardown_pending(&self, model_id: &str) -> bool {
963        let state = self
964            .state
965            .lock()
966            .unwrap_or_else(std::sync::PoisonError::into_inner);
967        let model_id = Self::resolve_model_identity(&state, model_id);
968        state.pending_teardown_models.contains_key(&model_id)
969    }
970
971    pub fn is_resident(&self, model_id: &str) -> bool {
972        let state = self
973            .state
974            .lock()
975            .unwrap_or_else(std::sync::PoisonError::into_inner);
976        let model_id = Self::resolve_model_identity(&state, model_id);
977        state
978            .resident_models
979            .iter()
980            .any(|(allocation_id, resident)| {
981                allocation_id == &model_id || resident.logical_model_id == model_id
982            })
983    }
984
985    pub fn resident_model_mb(&self) -> u64 {
986        let machine = self
987            .machine_ledger
988            .lock()
989            .unwrap_or_else(std::sync::PoisonError::into_inner);
990        machine
991            .resident_models
992            .values()
993            .chain(machine.pending_allocations.values())
994            .map(|resident| resident.weights_mb)
995            .fold(0, u64::saturating_add)
996    }
997
998    /// Test support for asserting one coordinator's accounting without
999    /// observing unrelated coordinators running in parallel.
1000    #[doc(hidden)]
1001    pub fn owned_resident_model_mb_for_testing(&self) -> u64 {
1002        let machine = self
1003            .machine_ledger
1004            .lock()
1005            .unwrap_or_else(std::sync::PoisonError::into_inner);
1006        machine
1007            .resident_models
1008            .iter()
1009            .chain(machine.pending_allocations.iter())
1010            .filter(|((owner_id, _), _)| *owner_id == self.owner_id)
1011            .map(|(_, resident)| resident.weights_mb)
1012            .fold(0, u64::saturating_add)
1013    }
1014
1015    pub fn active_request_count(&self, model_id: &str) -> usize {
1016        let state = self
1017            .state
1018            .lock()
1019            .unwrap_or_else(std::sync::PoisonError::into_inner);
1020        let model_id = Self::resolve_model_identity(&state, model_id);
1021        state
1022            .active_by_model
1023            .get(&model_id)
1024            .copied()
1025            .unwrap_or_default()
1026    }
1027
1028    /// Exact resident allocation keys for a logical model. Engine caches use
1029    /// the logical ID directly; independent voice/provider instances use
1030    /// allocation-qualified keys while retaining the same logical alias.
1031    pub fn resident_allocation_ids(&self, model_id: &str) -> Vec<String> {
1032        let state = self
1033            .state
1034            .lock()
1035            .unwrap_or_else(std::sync::PoisonError::into_inner);
1036        let model_id = Self::resolve_model_identity(&state, model_id);
1037        let mut allocations = state
1038            .resident_models
1039            .iter()
1040            .filter(|(allocation_id, resident)| {
1041                allocation_id.as_str() == model_id || resident.logical_model_id == model_id
1042            })
1043            .map(|(allocation_id, _)| allocation_id.clone())
1044            .collect::<Vec<_>>();
1045        if let Some(pending) = state.pending_teardown_models.get(&model_id) {
1046            allocations.extend(pending.iter().cloned());
1047        }
1048        allocations.sort();
1049        allocations.dedup();
1050        allocations
1051    }
1052
1053    fn default_weight_placement(&self) -> WeightPlacement {
1054        if matches!(self.hardware.gpu_backend, GpuBackend::Cuda)
1055            && self.hardware.gpu_memory_mb.is_some()
1056        {
1057            WeightPlacement::Accelerator
1058        } else {
1059            WeightPlacement::Host
1060        }
1061    }
1062
1063    fn mark_resident_with_placement(
1064        &self,
1065        logical_model_id: &str,
1066        allocation_id: &str,
1067        weights_mb: u64,
1068        placement: WeightPlacement,
1069    ) {
1070        let mut state = self
1071            .state
1072            .lock()
1073            .unwrap_or_else(std::sync::PoisonError::into_inner);
1074        let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
1075        let allocation = ResidentAllocation {
1076            weights_mb,
1077            placement,
1078            logical_model_id: logical_model_id.clone(),
1079        };
1080        state
1081            .resident_models
1082            .entry(allocation_id.to_string())
1083            .and_modify(|resident| {
1084                resident.weights_mb = resident.weights_mb.max(weights_mb);
1085                resident.placement = placement;
1086                resident.logical_model_id = logical_model_id.clone();
1087            })
1088            .or_insert(ResidentAllocation {
1089                weights_mb,
1090                placement,
1091                logical_model_id,
1092            });
1093        // Preserve the global lock order: scoped state, then machine ledger.
1094        self.machine_ledger
1095            .lock()
1096            .unwrap_or_else(std::sync::PoisonError::into_inner)
1097            .resident_models
1098            .insert((self.owner_id, allocation_id.to_string()), allocation);
1099    }
1100
1101    pub fn preflight(&self, model: &ModelSchema, context_tokens: usize) -> LocalLoadPreflight {
1102        let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
1103        let state = self
1104            .state
1105            .lock()
1106            .unwrap_or_else(std::sync::PoisonError::into_inner);
1107        let machine = self
1108            .machine_ledger
1109            .lock()
1110            .unwrap_or_else(std::sync::PoisonError::into_inner);
1111        self.preflight_locked(model, context_tokens, &state, &machine, live_available_mb)
1112    }
1113
1114    pub fn reserve(
1115        self: &Arc<Self>,
1116        model: &ModelSchema,
1117        context_tokens: usize,
1118    ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1119        let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
1120        let mut state = self
1121            .state
1122            .lock()
1123            .unwrap_or_else(std::sync::PoisonError::into_inner);
1124        let mut machine = self
1125            .machine_ledger
1126            .lock()
1127            .unwrap_or_else(std::sync::PoisonError::into_inner);
1128        let model_id = Self::resolve_model_identity(&state, &model.id);
1129        let estimate = self.model_memory_estimate(model, context_tokens);
1130        let mut preflight = self.preflight_estimate_locked(
1131            &model_id,
1132            estimate,
1133            self.default_weight_placement(),
1134            &state,
1135            &machine,
1136            live_available_mb,
1137        );
1138        if state.pending_teardown_models.contains_key(&model_id) {
1139            preflight.verdict = LocalLoadVerdict::PendingTeardown;
1140            return Err(LocalAdmissionError { preflight });
1141        }
1142        if state.maintenance_models.contains(&model_id) {
1143            preflight.verdict = LocalLoadVerdict::ModelMaintenance;
1144            return Err(LocalAdmissionError { preflight });
1145        }
1146        if !preflight.verdict.permits_static_fallback() {
1147            return Err(LocalAdmissionError { preflight });
1148        }
1149        if preflight.verdict == LocalLoadVerdict::LiveMemoryUnknown {
1150            tracing::warn!(
1151                model = %model.id,
1152                configured_ceiling_mb = preflight.configured_ceiling_mb,
1153                estimated_incremental_mb = preflight.estimated_incremental_mb,
1154                "live memory is unknown; proceeding under the static configured ceiling only"
1155            );
1156        }
1157        state.active_host_reservations_mb = state
1158            .active_host_reservations_mb
1159            .saturating_add(preflight.estimated_incremental_mb);
1160        state.active_accelerator_reservations_mb = state
1161            .active_accelerator_reservations_mb
1162            .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1163        let host = machine
1164            .active_host_by_owner
1165            .entry(self.owner_id)
1166            .or_default();
1167        *host = host.saturating_add(preflight.estimated_incremental_mb);
1168        let accelerator = machine
1169            .active_accelerator_by_owner
1170            .entry(self.owner_id)
1171            .or_default();
1172        *accelerator =
1173            accelerator.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1174        *state.active_by_model.entry(model_id.clone()).or_default() += 1;
1175        state.next_request_id = state.next_request_id.wrapping_add(1);
1176        let request_id = state.next_request_id;
1177        Ok(LocalLoadReservation {
1178            request_id,
1179            model_id: model_id.clone(),
1180            maintenance_model_id: model_id.clone(),
1181            weights_mb: preflight.estimate.weights_mb,
1182            reserved_incremental_mb: preflight.estimated_incremental_mb,
1183            reserved_accelerator_mb: preflight.accelerator_incremental_mb.unwrap_or_default(),
1184            cold_weights_reserved: !Self::has_resident_model(&state, &model_id),
1185            placement: self.default_weight_placement(),
1186            coordinator: Arc::clone(self),
1187            charge: Arc::new(ReservationCharge::new(
1188                model_id,
1189                preflight.estimated_incremental_mb,
1190                preflight.accelerator_incremental_mb.unwrap_or_default(),
1191                Arc::clone(self),
1192            )),
1193        })
1194    }
1195
1196    /// Reserve an installed host-memory allocation (Whisper/ONNX voice paths).
1197    /// The caller downloads first, measures the resulting artifact, then calls
1198    /// this immediately before constructing the native runtime.
1199    pub fn reserve_measured_host(
1200        self: &Arc<Self>,
1201        model_id: &str,
1202        measured_weights_bytes: u64,
1203        request_overhead_mb: u64,
1204    ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1205        self.reserve_measured(
1206            model_id,
1207            model_id,
1208            measured_weights_bytes,
1209            request_overhead_mb,
1210            WeightPlacement::Host,
1211        )
1212    }
1213
1214    pub fn reserve_measured_host_allocation(
1215        self: &Arc<Self>,
1216        logical_model_id: &str,
1217        allocation_id: &str,
1218        measured_weights_bytes: u64,
1219        request_overhead_mb: u64,
1220    ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1221        self.reserve_measured(
1222            logical_model_id,
1223            allocation_id,
1224            measured_weights_bytes,
1225            request_overhead_mb,
1226            WeightPlacement::Host,
1227        )
1228    }
1229
1230    fn reserve_measured(
1231        self: &Arc<Self>,
1232        logical_model_id: &str,
1233        allocation_id: &str,
1234        measured_weights_bytes: u64,
1235        request_overhead_mb: u64,
1236        placement: WeightPlacement,
1237    ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1238        let weights_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1239        let estimate = ModelMemoryEstimate {
1240            weights_mb,
1241            runtime_overhead_mb: request_overhead_mb,
1242            context_overhead_mb: 0,
1243            transient_margin_mb: 0,
1244            estimated_peak_mb: weights_mb.saturating_add(request_overhead_mb),
1245            evidence: ModelResourceEvidence::FileSystemMeasured,
1246        };
1247        let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
1248        let mut state = self
1249            .state
1250            .lock()
1251            .unwrap_or_else(std::sync::PoisonError::into_inner);
1252        let mut machine = self
1253            .machine_ledger
1254            .lock()
1255            .unwrap_or_else(std::sync::PoisonError::into_inner);
1256        let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
1257        let mut preflight = self.preflight_estimate_locked(
1258            allocation_id,
1259            estimate,
1260            placement,
1261            &state,
1262            &machine,
1263            live_available_mb,
1264        );
1265        if state
1266            .pending_teardown_models
1267            .contains_key(&logical_model_id)
1268        {
1269            preflight.verdict = LocalLoadVerdict::PendingTeardown;
1270        }
1271        if state.maintenance_models.contains(&logical_model_id) {
1272            preflight.verdict = LocalLoadVerdict::ModelMaintenance;
1273        }
1274        if !preflight.verdict.permits_static_fallback() {
1275            return Err(LocalAdmissionError { preflight });
1276        }
1277        state.active_host_reservations_mb = state
1278            .active_host_reservations_mb
1279            .saturating_add(preflight.estimated_incremental_mb);
1280        state.active_accelerator_reservations_mb = state
1281            .active_accelerator_reservations_mb
1282            .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1283        let host = machine
1284            .active_host_by_owner
1285            .entry(self.owner_id)
1286            .or_default();
1287        *host = host.saturating_add(preflight.estimated_incremental_mb);
1288        let accelerator = machine
1289            .active_accelerator_by_owner
1290            .entry(self.owner_id)
1291            .or_default();
1292        *accelerator =
1293            accelerator.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1294        *state
1295            .active_by_model
1296            .entry(logical_model_id.clone())
1297            .or_default() += 1;
1298        state.next_request_id = state.next_request_id.wrapping_add(1);
1299        Ok(LocalLoadReservation {
1300            request_id: state.next_request_id,
1301            model_id: allocation_id.to_string(),
1302            maintenance_model_id: logical_model_id.clone(),
1303            weights_mb,
1304            reserved_incremental_mb: preflight.estimated_incremental_mb,
1305            reserved_accelerator_mb: preflight.accelerator_incremental_mb.unwrap_or_default(),
1306            cold_weights_reserved: !state.resident_models.contains_key(allocation_id),
1307            placement,
1308            coordinator: Arc::clone(self),
1309            charge: Arc::new(ReservationCharge::new(
1310                logical_model_id,
1311                preflight.estimated_incremental_mb,
1312                preflight.accelerator_incremental_mb.unwrap_or_default(),
1313                Arc::clone(self),
1314            )),
1315        })
1316    }
1317
1318    /// Exclude new reservations for one model while removal/maintenance checks
1319    /// cache and cross-process lease state. The check and exclusion happen
1320    /// under the same mutex as `reserve`, closing the removal-vs-load race.
1321    pub fn begin_model_maintenance(
1322        self: &Arc<Self>,
1323        model_id: &str,
1324    ) -> Result<LocalModelMaintenanceGuard, ModelMaintenanceError> {
1325        let mut state = self
1326            .state
1327            .lock()
1328            .unwrap_or_else(std::sync::PoisonError::into_inner);
1329        let model_id = Self::resolve_model_identity(&state, model_id);
1330        if state
1331            .active_by_model
1332            .get(&model_id)
1333            .copied()
1334            .unwrap_or_default()
1335            > 0
1336        {
1337            return Err(ModelMaintenanceError::ModelInUse(model_id));
1338        }
1339        if !state.maintenance_models.insert(model_id.clone()) {
1340            return Err(ModelMaintenanceError::AlreadyInMaintenance(model_id));
1341        }
1342        Ok(LocalModelMaintenanceGuard {
1343            model_id,
1344            coordinator: Arc::clone(self),
1345        })
1346    }
1347
1348    fn preflight_locked(
1349        &self,
1350        model: &ModelSchema,
1351        context_tokens: usize,
1352        state: &AdmissionState,
1353        machine: &MachineAdmissionLedger,
1354        live_available_mb: Option<u64>,
1355    ) -> LocalLoadPreflight {
1356        let estimate = self.model_memory_estimate(model, context_tokens);
1357        let model_id = Self::resolve_model_identity(state, &model.id);
1358        self.preflight_estimate_locked(
1359            &model_id,
1360            estimate,
1361            self.default_weight_placement(),
1362            state,
1363            machine,
1364            live_available_mb,
1365        )
1366    }
1367
1368    fn preflight_estimate_locked(
1369        &self,
1370        model_id: &str,
1371        estimate: ModelMemoryEstimate,
1372        placement: WeightPlacement,
1373        state: &AdmissionState,
1374        machine: &MachineAdmissionLedger,
1375        live_available_mb: Option<u64>,
1376    ) -> LocalLoadPreflight {
1377        let policy = self.policy();
1378        let budget = policy.effective_budget(self.hardware.total_ram_mb);
1379        let host_resident_mb = machine
1380            .resident_models
1381            .values()
1382            .chain(machine.pending_allocations.values())
1383            .filter(|resident| resident.placement == WeightPlacement::Host)
1384            .map(|resident| resident.weights_mb)
1385            .fold(0, u64::saturating_add);
1386        let accelerator_resident_weights_mb = machine
1387            .resident_models
1388            .values()
1389            .chain(machine.pending_allocations.values())
1390            .filter(|resident| resident.placement == WeightPlacement::Accelerator)
1391            .map(|resident| resident.weights_mb)
1392            .fold(0, u64::saturating_add);
1393        let already_resident = Self::has_resident_model(state, model_id);
1394        let request_overhead_mb = estimate
1395            .context_overhead_mb
1396            .saturating_add(estimate.runtime_overhead_mb)
1397            .saturating_add(estimate.transient_margin_mb);
1398        let cold_weights_mb = if already_resident {
1399            0
1400        } else {
1401            estimate.weights_mb
1402        };
1403        let resident_model_mb = host_resident_mb;
1404        let host_cold_weights_mb = if placement == WeightPlacement::Host {
1405            cold_weights_mb
1406        } else {
1407            0
1408        };
1409        let estimated_incremental_mb = request_overhead_mb.saturating_add(host_cold_weights_mb);
1410        let accelerator_total_mb = (placement == WeightPlacement::Accelerator)
1411            .then_some(self.hardware.gpu_memory_mb)
1412            .flatten();
1413        let accelerator_resident_mb =
1414            (placement == WeightPlacement::Accelerator).then_some(accelerator_resident_weights_mb);
1415        let accelerator_incremental_mb =
1416            (placement == WeightPlacement::Accelerator).then_some(cold_weights_mb);
1417        let projected_static_mb = resident_model_mb
1418            .saturating_add(machine.active_host_by_owner.values().copied().sum::<u64>())
1419            .saturating_add(estimated_incremental_mb);
1420        let projected_accelerator_mb = accelerator_resident_weights_mb
1421            .saturating_add(
1422                machine
1423                    .active_accelerator_by_owner
1424                    .values()
1425                    .copied()
1426                    .sum::<u64>(),
1427            )
1428            .saturating_add(cold_weights_mb);
1429
1430        let verdict = if budget.effective_new_load_ceiling_mb == 0 && !already_resident {
1431            LocalLoadVerdict::DisabledByPolicy
1432        } else if (budget.configured_model_ceiling_mb != 0
1433            && projected_static_mb > budget.configured_model_ceiling_mb)
1434            || accelerator_total_mb.is_some_and(|vram_mb| projected_accelerator_mb > vram_mb)
1435        {
1436            LocalLoadVerdict::ExceedsConfiguredCeiling
1437        } else if let Some(available_mb) = live_available_mb {
1438            let unreserved_available_mb = available_mb
1439                .saturating_sub(machine.active_host_by_owner.values().copied().sum::<u64>());
1440            if unreserved_available_mb
1441                < estimated_incremental_mb.saturating_add(budget.emergency_reserve_mb)
1442            {
1443                LocalLoadVerdict::InsufficientLiveMemory
1444            } else {
1445                LocalLoadVerdict::Allowed
1446            }
1447        } else {
1448            LocalLoadVerdict::LiveMemoryUnknown
1449        };
1450
1451        LocalLoadPreflight {
1452            model_id: model_id.to_string(),
1453            estimate,
1454            configured_ceiling_mb: budget.configured_model_ceiling_mb,
1455            resident_model_mb,
1456            active_reservations_mb: machine.active_host_by_owner.values().copied().sum(),
1457            estimated_incremental_mb,
1458            accelerator_total_mb,
1459            accelerator_resident_mb,
1460            accelerator_incremental_mb,
1461            live_available_mb,
1462            emergency_reserve_mb: budget.emergency_reserve_mb,
1463            verdict,
1464        }
1465    }
1466
1467    fn has_resident_model(state: &AdmissionState, model_id: &str) -> bool {
1468        state
1469            .resident_models
1470            .iter()
1471            .any(|(allocation_id, resident)| {
1472                allocation_id == model_id || resident.logical_model_id == model_id
1473            })
1474    }
1475}
1476
1477#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
1478pub struct LocalAdmissionError {
1479    pub preflight: LocalLoadPreflight,
1480}
1481
1482#[derive(Clone, Debug, Error, PartialEq, Eq)]
1483pub enum ModelMaintenanceError {
1484    #[error("local model '{0}' is in use")]
1485    ModelInUse(String),
1486    #[error("local model '{0}' already has maintenance in progress")]
1487    AlreadyInMaintenance(String),
1488    #[error("failed to release local model residency: {0}")]
1489    ReleaseFailed(String),
1490    #[error("local worker did not acknowledge release of model '{0}'")]
1491    WorkerReleaseUnacknowledged(String),
1492    #[error("supervised local process did not acknowledge release of model '{0}'")]
1493    ProcessReleaseUnacknowledged(String),
1494    #[error("an in-process cache still has active work for local model '{0}'")]
1495    CacheReleaseBlocked(String),
1496    #[error("local model '{model_id}' still has resident allocations: {allocation_ids:?}")]
1497    ResidualResidency {
1498        model_id: String,
1499        allocation_ids: Vec<String>,
1500    },
1501}
1502
1503pub struct LocalModelMaintenanceGuard {
1504    model_id: String,
1505    coordinator: Arc<LocalAdmissionCoordinator>,
1506}
1507
1508impl Drop for LocalModelMaintenanceGuard {
1509    fn drop(&mut self) {
1510        self.coordinator
1511            .state
1512            .lock()
1513            .unwrap_or_else(std::sync::PoisonError::into_inner)
1514            .maintenance_models
1515            .remove(&self.model_id);
1516    }
1517}
1518
1519impl std::fmt::Display for LocalAdmissionError {
1520    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1521        write!(
1522            formatter,
1523            "local model '{}' blocked by {:?}: needs {} MB incremental, {} MB live available",
1524            self.preflight.model_id,
1525            self.preflight.verdict,
1526            self.preflight.estimated_incremental_mb,
1527            self.preflight
1528                .live_available_mb
1529                .map(|value| value.to_string())
1530                .unwrap_or_else(|| "unknown".into())
1531        )
1532    }
1533}
1534
1535impl std::error::Error for LocalAdmissionError {}
1536
1537/// RAII guard for one request's estimated incremental allocation. Dropping on
1538/// success, error, panic unwind, or cancellation releases the reservation.
1539pub struct LocalLoadReservation {
1540    request_id: u64,
1541    model_id: String,
1542    maintenance_model_id: String,
1543    weights_mb: u64,
1544    reserved_incremental_mb: u64,
1545    reserved_accelerator_mb: u64,
1546    cold_weights_reserved: bool,
1547    placement: WeightPlacement,
1548    coordinator: Arc<LocalAdmissionCoordinator>,
1549    charge: Arc<ReservationCharge>,
1550}
1551
1552/// Cloneable ownership for native work that outlives its awaiting request.
1553/// `spawn_blocking` tasks and supervised subprocess cleanup retain this lease
1554/// after cancellation so machine headroom is not advertised while allocation
1555/// work is still running.
1556#[derive(Clone)]
1557pub struct DetachedLocalLease {
1558    _charge: Arc<ReservationCharge>,
1559}
1560
1561struct ReservationCharge {
1562    maintenance_model_id: String,
1563    host_mb: AtomicU64,
1564    accelerator_mb: AtomicU64,
1565    cold_weights_transferred: AtomicBool,
1566    coordinator: Arc<LocalAdmissionCoordinator>,
1567    activity_lease: Mutex<Option<Arc<crate::model_management::ModelLease>>>,
1568}
1569
1570impl ReservationCharge {
1571    fn new(
1572        maintenance_model_id: String,
1573        host_mb: u64,
1574        accelerator_mb: u64,
1575        coordinator: Arc<LocalAdmissionCoordinator>,
1576    ) -> Self {
1577        Self {
1578            maintenance_model_id,
1579            host_mb: AtomicU64::new(host_mb),
1580            accelerator_mb: AtomicU64::new(accelerator_mb),
1581            cold_weights_transferred: AtomicBool::new(false),
1582            coordinator,
1583            activity_lease: Mutex::new(None),
1584        }
1585    }
1586
1587    fn update(&self, host_mb: u64, accelerator_mb: u64) {
1588        self.host_mb.store(host_mb, Ordering::Release);
1589        self.accelerator_mb.store(accelerator_mb, Ordering::Release);
1590    }
1591}
1592
1593impl Drop for ReservationCharge {
1594    fn drop(&mut self) {
1595        let host_mb = self.host_mb.load(Ordering::Acquire);
1596        let accelerator_mb = self.accelerator_mb.load(Ordering::Acquire);
1597        let mut state = self
1598            .coordinator
1599            .state
1600            .lock()
1601            .unwrap_or_else(std::sync::PoisonError::into_inner);
1602        let mut machine = self
1603            .coordinator
1604            .machine_ledger
1605            .lock()
1606            .unwrap_or_else(std::sync::PoisonError::into_inner);
1607        state.active_host_reservations_mb =
1608            state.active_host_reservations_mb.saturating_sub(host_mb);
1609        state.active_accelerator_reservations_mb = state
1610            .active_accelerator_reservations_mb
1611            .saturating_sub(accelerator_mb);
1612        if let Some(host) = machine
1613            .active_host_by_owner
1614            .get_mut(&self.coordinator.owner_id)
1615        {
1616            *host = host.saturating_sub(host_mb);
1617        }
1618        if let Some(accelerator) = machine
1619            .active_accelerator_by_owner
1620            .get_mut(&self.coordinator.owner_id)
1621        {
1622            *accelerator = accelerator.saturating_sub(accelerator_mb);
1623        }
1624        let maintenance_model_id =
1625            LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
1626        if let Some(active) = state.active_by_model.get_mut(&maintenance_model_id) {
1627            *active = active.saturating_sub(1);
1628            if *active == 0 {
1629                state.active_by_model.remove(&maintenance_model_id);
1630            }
1631        }
1632    }
1633}
1634
1635impl std::fmt::Debug for LocalLoadReservation {
1636    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1637        formatter
1638            .debug_struct("LocalLoadReservation")
1639            .field("request_id", &self.request_id)
1640            .field("model_id", &self.model_id)
1641            .field("reserved_incremental_mb", &self.reserved_incremental_mb)
1642            .field("reserved_accelerator_mb", &self.reserved_accelerator_mb)
1643            .finish_non_exhaustive()
1644    }
1645}
1646
1647impl LocalLoadReservation {
1648    pub fn request_id(&self) -> u64 {
1649        self.request_id
1650    }
1651
1652    /// Model whose allocation this reservation protects. Callers that walk an
1653    /// adaptive route can use this to avoid reusing one candidate's reservation
1654    /// for a different local fallback.
1655    pub fn model_id(&self) -> &str {
1656        &self.model_id
1657    }
1658
1659    pub fn authorizes_model(&self, model_id: &str) -> bool {
1660        self.model_id == model_id || self.maintenance_model_id == model_id
1661    }
1662
1663    /// Bind an admitted cache load to its concrete allocation owner. Distinct
1664    /// cache/runtime generations for one logical model must not peer-discount
1665    /// each other's physical weights.
1666    pub(crate) fn bind_allocation_id(&mut self, allocation_id: &str) {
1667        self.model_id = allocation_id.to_string();
1668    }
1669
1670    pub fn reserved_incremental_mb(&self) -> u64 {
1671        self.reserved_incremental_mb
1672    }
1673
1674    pub fn reconciled_weights_bytes(&self) -> u64 {
1675        self.weights_mb.saturating_mul(1024 * 1024)
1676    }
1677
1678    pub fn detached_lease(&self) -> DetachedLocalLease {
1679        DetachedLocalLease {
1680            _charge: self.charge.clone(),
1681        }
1682    }
1683
1684    pub(crate) fn attach_activity_lease(&mut self, lease: crate::model_management::ModelLease) {
1685        *self
1686            .charge
1687            .activity_lease
1688            .lock()
1689            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::new(lease));
1690    }
1691
1692    /// Atomically replace this request's cold-weight charge with a pending
1693    /// process allocation. Request overhead remains active until the request
1694    /// reservation and any detached cleanup lease are released.
1695    pub(crate) fn transfer_cold_weights_to_pending_allocation(
1696        &self,
1697        allocation_id: &str,
1698        measured_weights_bytes: u64,
1699    ) {
1700        let mut state = self
1701            .coordinator
1702            .state
1703            .lock()
1704            .unwrap_or_else(std::sync::PoisonError::into_inner);
1705        let logical_model_id =
1706            LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
1707        state
1708            .pending_teardown_models
1709            .entry(logical_model_id.clone())
1710            .or_default()
1711            .insert(allocation_id.to_string());
1712        let mut machine = self
1713            .coordinator
1714            .machine_ledger
1715            .lock()
1716            .unwrap_or_else(std::sync::PoisonError::into_inner);
1717        let key = (self.coordinator.owner_id, allocation_id.to_string());
1718        let machine_resident = machine.resident_models.remove(&key);
1719        let scoped_resident = state.resident_models.get(allocation_id).cloned();
1720        let resident = machine_resident.or(scoped_resident);
1721        let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1722        machine
1723            .pending_allocations
1724            .entry(key)
1725            .and_modify(|pending| {
1726                pending.weights_mb = pending.weights_mb.max(measured_mb);
1727                pending.logical_model_id = logical_model_id.clone();
1728            })
1729            .or_insert_with(|| {
1730                resident.unwrap_or(ResidentAllocation {
1731                    weights_mb: measured_mb,
1732                    placement: self.placement,
1733                    logical_model_id,
1734                })
1735            });
1736
1737        if self.cold_weights_reserved
1738            && self
1739                .charge
1740                .cold_weights_transferred
1741                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
1742                .is_ok()
1743        {
1744            match self.placement {
1745                WeightPlacement::Host => {
1746                    state.active_host_reservations_mb = state
1747                        .active_host_reservations_mb
1748                        .saturating_sub(self.weights_mb);
1749                    let active = machine
1750                        .active_host_by_owner
1751                        .entry(self.coordinator.owner_id)
1752                        .or_default();
1753                    *active = active.saturating_sub(self.weights_mb);
1754                    let _ = self.charge.host_mb.fetch_update(
1755                        Ordering::AcqRel,
1756                        Ordering::Acquire,
1757                        |host_mb| Some(host_mb.saturating_sub(self.weights_mb)),
1758                    );
1759                }
1760                WeightPlacement::Accelerator => {
1761                    state.active_accelerator_reservations_mb = state
1762                        .active_accelerator_reservations_mb
1763                        .saturating_sub(self.weights_mb);
1764                    let active = machine
1765                        .active_accelerator_by_owner
1766                        .entry(self.coordinator.owner_id)
1767                        .or_default();
1768                    *active = active.saturating_sub(self.weights_mb);
1769                    let _ = self.charge.accelerator_mb.fetch_update(
1770                        Ordering::AcqRel,
1771                        Ordering::Acquire,
1772                        |accelerator_mb| Some(accelerator_mb.saturating_sub(self.weights_mb)),
1773                    );
1774                }
1775            }
1776        }
1777    }
1778
1779    fn sync_shared_charge(&self) {
1780        self.charge
1781            .update(self.reserved_incremental_mb, self.reserved_accelerator_mb);
1782    }
1783
1784    /// Replace the catalog estimate with an installed on-disk measurement.
1785    /// Check and reservation replacement happen under the same coordinator
1786    /// mutex; rejection leaves the original reservation untouched.
1787    pub fn reconcile_measured_weights(
1788        &mut self,
1789        measured_weights_bytes: u64,
1790    ) -> Result<LocalLoadPreflight, LocalAdmissionError> {
1791        if !self.cold_weights_reserved {
1792            let live_available_mb = self
1793                .coordinator
1794                .live_probe
1795                .available_memory_mb()
1796                .ok()
1797                .flatten();
1798            let state = self
1799                .coordinator
1800                .state
1801                .lock()
1802                .unwrap_or_else(std::sync::PoisonError::into_inner);
1803            let machine = self
1804                .coordinator
1805                .machine_ledger
1806                .lock()
1807                .unwrap_or_else(std::sync::PoisonError::into_inner);
1808            if LocalAdmissionCoordinator::has_resident_model(&state, &self.model_id) {
1809                let estimate = self.measured_estimate(self.weights_mb);
1810                return Ok(self.coordinator.preflight_estimate_locked(
1811                    &self.model_id,
1812                    estimate,
1813                    self.placement,
1814                    &state,
1815                    &machine,
1816                    live_available_mb,
1817                ));
1818            }
1819            drop(state);
1820            self.cold_weights_reserved = true;
1821        }
1822        let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1823        let live_available_mb = self
1824            .coordinator
1825            .live_probe
1826            .available_memory_mb()
1827            .ok()
1828            .flatten();
1829        let mut state = self
1830            .coordinator
1831            .state
1832            .lock()
1833            .unwrap_or_else(std::sync::PoisonError::into_inner);
1834        let mut machine = self
1835            .coordinator
1836            .machine_ledger
1837            .lock()
1838            .unwrap_or_else(std::sync::PoisonError::into_inner);
1839        let published_by_peer =
1840            LocalAdmissionCoordinator::has_resident_model(&state, &self.model_id);
1841        let mut without_current = state.clone();
1842        without_current.active_host_reservations_mb = without_current
1843            .active_host_reservations_mb
1844            .saturating_sub(self.reserved_incremental_mb);
1845        without_current.active_accelerator_reservations_mb = without_current
1846            .active_accelerator_reservations_mb
1847            .saturating_sub(self.reserved_accelerator_mb);
1848        let mut without_machine = machine.clone();
1849        let host = without_machine
1850            .active_host_by_owner
1851            .entry(self.coordinator.owner_id)
1852            .or_default();
1853        *host = host.saturating_sub(self.reserved_incremental_mb);
1854        let accelerator = without_machine
1855            .active_accelerator_by_owner
1856            .entry(self.coordinator.owner_id)
1857            .or_default();
1858        *accelerator = accelerator.saturating_sub(self.reserved_accelerator_mb);
1859        let estimate = self.measured_estimate(measured_mb);
1860        let preflight = self.coordinator.preflight_estimate_locked(
1861            &self.model_id,
1862            estimate,
1863            self.placement,
1864            &without_current,
1865            &without_machine,
1866            live_available_mb,
1867        );
1868        if !preflight.verdict.permits_static_fallback() {
1869            return Err(LocalAdmissionError { preflight });
1870        }
1871        state.active_host_reservations_mb = without_current
1872            .active_host_reservations_mb
1873            .saturating_add(preflight.estimated_incremental_mb);
1874        state.active_accelerator_reservations_mb = without_current
1875            .active_accelerator_reservations_mb
1876            .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1877        machine.active_host_by_owner.insert(
1878            self.coordinator.owner_id,
1879            without_machine
1880                .active_host_by_owner
1881                .get(&self.coordinator.owner_id)
1882                .copied()
1883                .unwrap_or_default()
1884                .saturating_add(preflight.estimated_incremental_mb),
1885        );
1886        machine.active_accelerator_by_owner.insert(
1887            self.coordinator.owner_id,
1888            without_machine
1889                .active_accelerator_by_owner
1890                .get(&self.coordinator.owner_id)
1891                .copied()
1892                .unwrap_or_default()
1893                .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default()),
1894        );
1895        self.weights_mb = measured_mb;
1896        self.reserved_incremental_mb = preflight.estimated_incremental_mb;
1897        self.reserved_accelerator_mb = preflight.accelerator_incremental_mb.unwrap_or_default();
1898        if published_by_peer {
1899            self.cold_weights_reserved = false;
1900        }
1901        self.sync_shared_charge();
1902        Ok(preflight)
1903    }
1904
1905    fn measured_estimate(&self, weights_mb: u64) -> ModelMemoryEstimate {
1906        let host_cold_mb = if self.cold_weights_reserved && self.placement == WeightPlacement::Host
1907        {
1908            self.weights_mb
1909        } else {
1910            0
1911        };
1912        let request_overhead_mb = self.reserved_incremental_mb.saturating_sub(host_cold_mb);
1913        ModelMemoryEstimate {
1914            weights_mb,
1915            runtime_overhead_mb: request_overhead_mb,
1916            context_overhead_mb: 0,
1917            transient_margin_mb: 0,
1918            estimated_peak_mb: weights_mb.saturating_add(request_overhead_mb),
1919            evidence: ModelResourceEvidence::FileSystemMeasured,
1920        }
1921    }
1922
1923    /// Atomically transfer the cold portion of this request reservation into
1924    /// resident accounting at the exact point a cache publishes the backend.
1925    pub fn publish_resident_weights(&mut self, measured_weights_bytes: u64) {
1926        let allocation_id = self.model_id.clone();
1927        self.publish_resident_weights_as(&allocation_id, measured_weights_bytes);
1928    }
1929
1930    /// Conservatively account an external allocation whose acknowledgement
1931    /// reported a different concrete allocation id than the requested model.
1932    /// This is only used when the external owner cannot confirm teardown.
1933    pub fn publish_resident_weights_as(
1934        &mut self,
1935        allocation_id: &str,
1936        measured_weights_bytes: u64,
1937    ) {
1938        // Publication happens after the allocation already exists. It must not
1939        // run another live-memory probe or reject at that point: doing so can
1940        // return an error while a cache/process remains resident but
1941        // unaccounted. The caller reconciles before loading; if an external
1942        // allocator reports a different final size, conservatively account the
1943        // actual size atomically even when it exceeds the earlier reservation.
1944        let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1945        let mut state = self
1946            .coordinator
1947            .state
1948            .lock()
1949            .unwrap_or_else(std::sync::PoisonError::into_inner);
1950        let mut machine = self
1951            .coordinator
1952            .machine_ledger
1953            .lock()
1954            .unwrap_or_else(std::sync::PoisonError::into_inner);
1955        let key = (self.coordinator.owner_id, allocation_id.to_string());
1956        let pending = machine.pending_allocations.remove(&key);
1957        let logical_model_id =
1958            LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
1959        if let Some(allocations) = state.pending_teardown_models.get_mut(&logical_model_id) {
1960            allocations.remove(allocation_id);
1961            if allocations.is_empty() {
1962                state.pending_teardown_models.remove(&logical_model_id);
1963            }
1964        }
1965        if !self.cold_weights_reserved && pending.is_none() {
1966            return;
1967        }
1968        if self.cold_weights_reserved {
1969            let cold_weights_transferred =
1970                self.charge.cold_weights_transferred.load(Ordering::Acquire);
1971            match self.placement {
1972                WeightPlacement::Host => {
1973                    if !cold_weights_transferred {
1974                        state.active_host_reservations_mb = state
1975                            .active_host_reservations_mb
1976                            .saturating_sub(self.weights_mb);
1977                        let active = machine
1978                            .active_host_by_owner
1979                            .entry(self.coordinator.owner_id)
1980                            .or_default();
1981                        *active = active.saturating_sub(self.weights_mb);
1982                    }
1983                    self.reserved_incremental_mb =
1984                        self.reserved_incremental_mb.saturating_sub(self.weights_mb);
1985                }
1986                WeightPlacement::Accelerator => {
1987                    if !cold_weights_transferred {
1988                        state.active_accelerator_reservations_mb = state
1989                            .active_accelerator_reservations_mb
1990                            .saturating_sub(self.weights_mb);
1991                        let active = machine
1992                            .active_accelerator_by_owner
1993                            .entry(self.coordinator.owner_id)
1994                            .or_default();
1995                        *active = active.saturating_sub(self.weights_mb);
1996                    }
1997                    self.reserved_accelerator_mb =
1998                        self.reserved_accelerator_mb.saturating_sub(self.weights_mb);
1999                }
2000            }
2001        }
2002        let resident = ResidentAllocation {
2003            weights_mb: pending
2004                .as_ref()
2005                .map(|allocation| allocation.weights_mb)
2006                .unwrap_or_default()
2007                .max(measured_mb),
2008            placement: pending
2009                .as_ref()
2010                .map(|allocation| allocation.placement)
2011                .unwrap_or(self.placement),
2012            logical_model_id,
2013        };
2014        state
2015            .resident_models
2016            .insert(allocation_id.to_string(), resident.clone());
2017        machine.resident_models.insert(key, resident.clone());
2018        if let Some(lease) = self
2019            .charge
2020            .activity_lease
2021            .lock()
2022            .unwrap_or_else(std::sync::PoisonError::into_inner)
2023            .clone()
2024        {
2025            self.coordinator
2026                .resident_activity_leases
2027                .lock()
2028                .unwrap_or_else(std::sync::PoisonError::into_inner)
2029                .insert(allocation_id.to_string(), lease);
2030        }
2031        self.weights_mb = resident.weights_mb;
2032        self.cold_weights_reserved = false;
2033        self.charge
2034            .cold_weights_transferred
2035            .store(false, Ordering::Release);
2036        self.sync_shared_charge();
2037    }
2038
2039    /// Transfer cold weight accounting into the resident set after a loader
2040    /// successfully publishes the backend into its cache.
2041    pub fn commit_resident_weights(&mut self) {
2042        self.publish_resident_weights(self.weights_mb * 1024 * 1024);
2043    }
2044}
2045
2046#[derive(Clone, Debug, PartialEq, Eq)]
2047struct KvGeometry {
2048    num_hidden_layers: usize,
2049    num_key_value_heads: usize,
2050    head_dim: usize,
2051    kv_bytes_per_element: usize,
2052    hidden_size: Option<usize>,
2053}
2054
2055/// Installed checkpoints' own KV geometry, keyed by model id: the input that
2056/// lets the recommendation estimate size an installed row's KV cache from its
2057/// `config.json`, as load-time admission does, instead of from the parameter
2058/// count. A row absent here — not installed, or a config without readable
2059/// geometry — keeps the parameter-count fallback, as it does at admission.
2060/// [`CheckpointGeometry::default`] holds nothing, so every row falls back.
2061#[derive(Clone, Debug, Default)]
2062pub struct CheckpointGeometry {
2063    by_id: HashMap<String, KvGeometry>,
2064}
2065
2066impl CheckpointGeometry {
2067    /// Read each row's `<models_dir>/<name>/config.json`, the file admission
2068    /// reads. Callers that hold an engine use
2069    /// [`LocalAdmissionCoordinator::checkpoint_geometry`] instead, which shares
2070    /// admission's models directory and cache.
2071    pub fn read<'a>(models: impl IntoIterator<Item = &'a ModelSchema>, models_dir: &Path) -> Self {
2072        Self {
2073            by_id: models
2074                .into_iter()
2075                .filter_map(|model| {
2076                    Some((model.id.clone(), kv_geometry_for(model, Some(models_dir))?))
2077                })
2078                .collect(),
2079        }
2080    }
2081
2082    fn get(&self, model: &ModelSchema) -> Option<&KvGeometry> {
2083        self.by_id.get(&model.id)
2084    }
2085}
2086
2087/// Where multimodal checkpoints nest their language model's settings:
2088/// `text_config` (gemma-4, Qwen-VL), `llm_config` (InternVL),
2089/// `language_config` (DeepSeek-VL). Defined here, not beside the MLX loader,
2090/// because the loader compiles only on Apple Silicon and this file everywhere.
2091pub(crate) const NESTED_TEXT_CONFIGS: [&str; 3] = ["text_config", "llm_config", "language_config"];
2092
2093#[derive(Deserialize)]
2094struct ModelGeometryConfig {
2095    num_hidden_layers: Option<usize>,
2096    num_key_value_heads: Option<usize>,
2097    head_dim: Option<usize>,
2098    hidden_size: Option<usize>,
2099    num_attention_heads: Option<usize>,
2100    kv_cache_dtype: Option<String>,
2101    torch_dtype: Option<String>,
2102}
2103
2104fn dtype_bytes(dtype: Option<&str>) -> usize {
2105    match dtype.map(str::to_ascii_lowercase).as_deref() {
2106        Some("float32" | "fp32" | "f32") => 4,
2107        Some("float64" | "fp64" | "f64") => 8,
2108        Some("float8" | "fp8" | "int8" | "uint8" | "i8" | "u8") => 1,
2109        // MLX-community quantizes weights, while its KV tensors remain fp16 or
2110        // bf16. Missing and unrecognized dtype declarations therefore retain
2111        // the safe two-byte default rather than inheriting weight bit width.
2112        _ => 2,
2113    }
2114}
2115
2116/// What identifies one version of a checkpoint's `config.json`: a remove, a
2117/// replace or a re-download changes it.
2118#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2119struct ConfigStamp {
2120    modified: Option<std::time::SystemTime>,
2121    len: u64,
2122}
2123
2124impl ConfigStamp {
2125    fn of(path: &Path) -> Option<Self> {
2126        let metadata = std::fs::metadata(path).ok()?;
2127        Some(Self {
2128            modified: metadata.modified().ok(),
2129            len: metadata.len(),
2130        })
2131    }
2132}
2133
2134fn kv_config_path(model: &ModelSchema, models_dir: &Path) -> PathBuf {
2135    models_dir.join(&model.name).join("config.json")
2136}
2137
2138/// Resolve KV geometry from `<models_dir>/<model.name>/config.json`.
2139///
2140/// The config may explicitly state `kv_cache_dtype` (or the model-wide
2141/// `torch_dtype`); otherwise KV elements default to two bytes for fp16/bf16.
2142/// Missing, unreadable, malformed, or incomplete configs have no geometry and
2143/// use the conservative parameter-count fallback below.
2144///
2145/// A multimodal checkpoint (gemma-4's `gemma4_unified`) keeps its language
2146/// model's geometry under a nested config, so when the top level has no layer
2147/// count, the nested one is read — the same keys the loader reads.
2148fn kv_geometry_for(model: &ModelSchema, models_dir: Option<&Path>) -> Option<KvGeometry> {
2149    let bytes = std::fs::read(kv_config_path(model, models_dir?)).ok()?;
2150    let raw: serde_json::Value = serde_json::from_slice(&bytes).ok()?;
2151    let source = if raw.get("num_hidden_layers").is_some() {
2152        &raw
2153    } else {
2154        NESTED_TEXT_CONFIGS
2155            .iter()
2156            .find_map(|key| {
2157                raw.get(*key)
2158                    .filter(|v| v.get("num_hidden_layers").is_some())
2159            })
2160            .unwrap_or(&raw)
2161    };
2162    let mut config: ModelGeometryConfig = serde_json::from_value(source.clone()).ok()?;
2163    // The KV dtype may be declared once, at the top.
2164    if config.kv_cache_dtype.is_none() && config.torch_dtype.is_none() {
2165        config.torch_dtype = raw
2166            .get("torch_dtype")
2167            .and_then(serde_json::Value::as_str)
2168            .map(str::to_string);
2169    }
2170    let num_hidden_layers = config.num_hidden_layers.filter(|value| *value > 0)?;
2171    let num_key_value_heads = config.num_key_value_heads.filter(|value| *value > 0)?;
2172    let head_dim = config.head_dim.filter(|value| *value > 0).or_else(|| {
2173        let hidden_size = config.hidden_size?;
2174        let attention_heads = config.num_attention_heads.filter(|value| *value > 0)?;
2175        hidden_size
2176            .is_multiple_of(attention_heads)
2177            .then_some(hidden_size / attention_heads)
2178    })?;
2179    let dtype = config
2180        .kv_cache_dtype
2181        .as_deref()
2182        .or(config.torch_dtype.as_deref());
2183    Some(KvGeometry {
2184        num_hidden_layers,
2185        num_key_value_heads,
2186        head_dim,
2187        kv_bytes_per_element: dtype_bytes(dtype),
2188        hidden_size: config.hidden_size.filter(|value| *value > 0),
2189    })
2190}
2191
2192/// Binary-megabyte KV allocation for decoder geometry and context length:
2193/// `2 × layers × KV heads × head dimension × element bytes × tokens`.
2194fn kv_cache_mb_for(geometry: &KvGeometry, context_tokens: usize) -> u64 {
2195    let bytes = 2_u128
2196        .saturating_mul(geometry.num_hidden_layers as u128)
2197        .saturating_mul(geometry.num_key_value_heads as u128)
2198        .saturating_mul(geometry.head_dim as u128)
2199        .saturating_mul(geometry.kv_bytes_per_element as u128)
2200        .saturating_mul(context_tokens as u128);
2201    bytes.div_ceil(BYTES_PER_MB).min(u64::MAX as u128) as u64
2202}
2203
2204/// Conservative unknown-geometry fallback. Qwen3-8B fp16 KV uses 147,456
2205/// bytes/token (`2 × 36 layers × 8 KV heads × 128 head_dim × 2 bytes`). Dense
2206/// 0.6B–14B models sit within roughly 0.8–1.2× of this figure per token; larger
2207/// active parameter counts scale it upward. Installed, known geometry wins.
2208fn fallback_kv_cache_mb(model: &ModelSchema, context_tokens: usize) -> u64 {
2209    let active_billions = model_parameter_billions_active(model) as f64;
2210    let bytes_per_token = 147_456.0 * (active_billions / 8.0).max(1.0);
2211    ((bytes_per_token * context_tokens as f64) / BYTES_PER_MB as f64).ceil() as u64
2212}
2213
2214/// Four fp16 `[tokens × hidden]` tensors conservatively alive together during
2215/// prompt prefill. This is additive to persistent KV allocation.
2216fn prefill_activation_mb(geometry: Option<&KvGeometry>, context_tokens: usize) -> u64 {
2217    let Some(hidden_size) = geometry.and_then(|geometry| geometry.hidden_size) else {
2218        return 0;
2219    };
2220    let bytes = (hidden_size as u128)
2221        .saturating_mul(context_tokens as u128)
2222        .saturating_mul(2)
2223        .saturating_mul(4);
2224    bytes.div_ceil(BYTES_PER_MB).min(u64::MAX as u128) as u64
2225}
2226
2227/// Central recommendation-time memory estimator.
2228pub fn estimate_model_memory(
2229    model: &ModelSchema,
2230    hardware: &HardwareInfo,
2231    context_tokens: usize,
2232) -> ModelMemoryEstimate {
2233    estimate_model_memory_with_measured_weights(model, hardware, context_tokens, None)
2234}
2235
2236/// The recommendation estimate for a row whose installed checkpoint may
2237/// declare its own KV geometry: the KV cache is sized from that geometry when
2238/// `geometry` holds the row, and from the parameter count otherwise.
2239pub fn estimate_model_memory_for_checkpoint(
2240    model: &ModelSchema,
2241    hardware: &HardwareInfo,
2242    context_tokens: usize,
2243    geometry: &CheckpointGeometry,
2244) -> ModelMemoryEstimate {
2245    estimate_model_memory_with_geometry(model, hardware, context_tokens, None, geometry.get(model))
2246}
2247
2248/// The same estimate with optional physical weight evidence. Registry/load
2249/// paths can provide a measured value without duplicating the remaining math.
2250pub fn estimate_model_memory_with_measured_weights(
2251    model: &ModelSchema,
2252    hardware: &HardwareInfo,
2253    context_tokens: usize,
2254    measured_weights_mb: Option<u64>,
2255) -> ModelMemoryEstimate {
2256    estimate_model_memory_with_geometry(model, hardware, context_tokens, measured_weights_mb, None)
2257}
2258
2259fn estimate_model_memory_with_geometry(
2260    model: &ModelSchema,
2261    hardware: &HardwareInfo,
2262    context_tokens: usize,
2263    measured_weights_mb: Option<u64>,
2264    geometry: Option<&KvGeometry>,
2265) -> ModelMemoryEstimate {
2266    let declared = model.cost.ram_mb.or(model.cost.size_mb);
2267    let (weights_mb, evidence) = if let Some(measured) = measured_weights_mb {
2268        (measured, ModelResourceEvidence::FileSystemMeasured)
2269    } else if let Some(declared) = declared {
2270        (
2271            declared.max(model.cost.size_mb.unwrap_or(0)),
2272            ModelResourceEvidence::CatalogExact,
2273        )
2274    } else {
2275        (
2276            heuristic_weights_mb(model),
2277            ModelResourceEvidence::Heuristic,
2278        )
2279    };
2280    let kv_cache_mb = geometry
2281        .map(|geometry| kv_cache_mb_for(geometry, context_tokens))
2282        .unwrap_or_else(|| fallback_kv_cache_mb(model, context_tokens));
2283    let context_overhead_mb =
2284        kv_cache_mb.saturating_add(prefill_activation_mb(geometry, context_tokens));
2285    let runtime_overhead_mb = backend_runtime_overhead_mb(hardware);
2286    let transient_margin_mb = TRANSIENT_ALLOCATION_MARGIN_MB;
2287    let estimated_peak_mb = weights_mb
2288        .saturating_add(context_overhead_mb)
2289        .saturating_add(runtime_overhead_mb)
2290        .saturating_add(transient_margin_mb);
2291
2292    ModelMemoryEstimate {
2293        weights_mb,
2294        runtime_overhead_mb,
2295        context_overhead_mb,
2296        transient_margin_mb,
2297        estimated_peak_mb,
2298        evidence,
2299    }
2300}
2301
2302fn backend_runtime_overhead_mb(hardware: &HardwareInfo) -> u64 {
2303    match hardware.gpu_backend {
2304        GpuBackend::Metal => METAL_RUNTIME_OVERHEAD_MB,
2305        GpuBackend::Cuda => CUDA_RUNTIME_OVERHEAD_MB,
2306        _ => CPU_RUNTIME_OVERHEAD_MB,
2307    }
2308}
2309
2310fn heuristic_weights_mb(model: &ModelSchema) -> u64 {
2311    let billions = model_parameter_billions_total(model);
2312    (billions as f64 * 600.0).ceil() as u64
2313}
2314
2315pub(crate) fn model_parameter_billions_active(model: &ModelSchema) -> f32 {
2316    model
2317        .param_count
2318        .split_once('(')
2319        .and_then(|(_, rest)| rest.split_once("active"))
2320        .and_then(|(number, _)| parse_parameter_billions(number))
2321        .unwrap_or_else(|| model_parameter_billions_total(model))
2322}
2323
2324pub(crate) fn model_parameter_billions_total(model: &ModelSchema) -> f32 {
2325    parse_parameter_billions(&model.param_count).unwrap_or_else(|| {
2326        let size_mb = model.size_mb();
2327        if size_mb > 0 {
2328            (size_mb as f32 / 600.0).max(0.1)
2329        } else {
2330            0.0
2331        }
2332    })
2333}
2334
2335pub(crate) fn parse_parameter_billions(value: &str) -> Option<f32> {
2336    let value = value.trim();
2337    let number: String = value
2338        .chars()
2339        .take_while(|character| character.is_ascii_digit() || *character == '.')
2340        .collect();
2341    let parsed: f32 = number.parse().ok()?;
2342    if value[number.len()..]
2343        .trim_start()
2344        .to_ascii_lowercase()
2345        .starts_with('m')
2346    {
2347        Some(parsed / 1_000.0)
2348    } else {
2349        Some(parsed)
2350    }
2351}
2352
2353/// Evaluate policy math for the hardware CAR actually detected.
2354pub fn evaluate_resources(hardware: &HardwareInfo, policy: &ResourcePolicy) -> ResourceEvaluation {
2355    let accelerator_memory = match (hardware.gpu_backend.clone(), hardware.gpu_memory_mb) {
2356        (GpuBackend::Cuda, Some(total_mb)) => Some(AcceleratorResourceBudget {
2357            total_mb,
2358            budget_mb: total_mb,
2359        }),
2360        _ => None,
2361    };
2362
2363    ResourceEvaluation {
2364        host_memory: policy.effective_budget(hardware.total_ram_mb),
2365        accelerator_memory,
2366    }
2367}
2368
2369fn percent_of(total_mb: u64, percent: u64) -> u64 {
2370    let value = (total_mb as u128).saturating_mul(percent as u128) / 100;
2371    value.min(u64::MAX as u128) as u64
2372}
2373
2374fn minimum_emergency_reserve(total_memory_mb: u64) -> u64 {
2375    MINIMUM_EMERGENCY_RESERVE_MB.max(percent_of(total_memory_mb, EMERGENCY_RESERVE_PERCENT))
2376}
2377
2378/// Persistence errors that callers can distinguish from safe defaulting on a
2379/// missing or malformed policy document.
2380#[derive(Debug, Error)]
2381pub enum ResourcePolicyError {
2382    #[error("resource policy I/O failed: {0}")]
2383    Io(#[from] io::Error),
2384    #[error("resource policy serialization failed: {0}")]
2385    Serialization(#[from] serde_json::Error),
2386    #[error("Custom model RAM must be finite, nonnegative, and a 0.5 GB increment; got {0}")]
2387    InvalidCustomGigabytes(f64),
2388    #[error("invalid resource policy: {reason}")]
2389    InvalidPolicy { reason: String },
2390}
2391
2392/// Where the effective policy returned by a repository load came from.
2393#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
2394#[serde(rename_all = "snake_case")]
2395pub enum ResourcePolicyLoadSource {
2396    Loaded,
2397    MissingDefault,
2398    CorruptDefault,
2399}
2400
2401/// A policy plus recoverable evidence about safe defaulting.
2402#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
2403pub struct ResourcePolicyLoadEvidence {
2404    pub policy: ResourcePolicy,
2405    pub source: ResourcePolicyLoadSource,
2406    pub warning: Option<String>,
2407}
2408
2409pub trait ResourcePolicyRepository: Send + Sync {
2410    fn load(&self) -> Result<ResourcePolicy, ResourcePolicyError>;
2411    fn save(&self, policy: &ResourcePolicy) -> Result<(), ResourcePolicyError>;
2412}
2413
2414/// Private, atomic resource-policy storage rooted at `CAR_HOME`.
2415#[derive(Clone, Debug)]
2416pub struct FileResourcePolicyRepository {
2417    root: PathBuf,
2418}
2419
2420impl Default for FileResourcePolicyRepository {
2421    fn default() -> Self {
2422        Self::new(car_home::root_or_relative())
2423    }
2424}
2425
2426/// File under the CAR state root recording that the saved resource policy
2427/// is the concierge's choice, not the user's.
2428pub const CONCIERGE_POLICY_MARKER: &str = "model-resource-policy.concierge.json";
2429
2430#[derive(Debug, Serialize, Deserialize)]
2431struct ConciergePolicyMarker {
2432    policy: ResourcePolicy,
2433    reason: String,
2434    at: u64,
2435    /// The concierge undid its own raise: it does not raise again.
2436    #[serde(default)]
2437    reverted: bool,
2438}
2439
2440/// What the concierge's marker records about the saved policy.
2441#[derive(Debug, Clone, PartialEq, Eq)]
2442pub struct ConciergeChoice {
2443    pub policy: ResourcePolicy,
2444    pub at: u64,
2445    pub reverted: bool,
2446}
2447
2448impl FileResourcePolicyRepository {
2449    pub fn new(root: PathBuf) -> Self {
2450        Self { root }
2451    }
2452
2453    /// The policy the concierge last saved, if the saved policy was its
2454    /// choice. A user's save clears this ([`Self::clear_concierge_marker`]),
2455    /// so the saved policy differing from it, or it being absent while a
2456    /// policy is saved, means the user chose.
2457    pub fn concierge_marker(&self) -> Option<ResourcePolicy> {
2458        self.concierge_choice().map(|choice| choice.policy)
2459    }
2460
2461    /// The full marker: the policy, when it was set, and whether it was a
2462    /// reversal of the concierge's own raise.
2463    pub fn concierge_choice(&self) -> Option<ConciergeChoice> {
2464        let bytes = std::fs::read(self.root.join(CONCIERGE_POLICY_MARKER)).ok()?;
2465        serde_json::from_slice::<ConciergePolicyMarker>(&bytes)
2466            .ok()
2467            .map(|marker| ConciergeChoice {
2468                policy: marker.policy,
2469                at: marker.at,
2470                reverted: marker.reverted,
2471            })
2472    }
2473
2474    /// Save `policy` as the concierge's choice: the policy, then the marker.
2475    pub fn save_concierge_choice(
2476        &self,
2477        policy: &ResourcePolicy,
2478        reason: &str,
2479        at: u64,
2480        reverted: bool,
2481    ) -> Result<(), ResourcePolicyError> {
2482        ResourcePolicyRepository::save(self, policy)?;
2483        let marker = ConciergePolicyMarker {
2484            policy: policy.clone(),
2485            reason: reason.to_string(),
2486            at,
2487            reverted,
2488        };
2489        let path = self.root.join(CONCIERGE_POLICY_MARKER);
2490        let tmp = self.root.join(format!(".{CONCIERGE_POLICY_MARKER}.tmp"));
2491        let bytes = serde_json::to_vec_pretty(&marker).map_err(io::Error::other)?;
2492        std::fs::write(&tmp, bytes)?;
2493        std::fs::rename(&tmp, &path)?;
2494        Ok(())
2495    }
2496
2497    /// The user chose: the saved policy is theirs from now on.
2498    pub fn clear_concierge_marker(&self) -> Result<(), ResourcePolicyError> {
2499        match std::fs::remove_file(self.root.join(CONCIERGE_POLICY_MARKER)) {
2500            Err(error) if error.kind() != io::ErrorKind::NotFound => Err(error.into()),
2501            _ => Ok(()),
2502        }
2503    }
2504
2505    pub fn path(&self) -> PathBuf {
2506        self.root.join(RESOURCE_POLICY_FILE)
2507    }
2508
2509    /// Load the effective policy without losing evidence that CAR recovered
2510    /// from a missing or invalid source document.
2511    pub fn load_with_evidence(&self) -> Result<ResourcePolicyLoadEvidence, ResourcePolicyError> {
2512        let path = self.path();
2513        let file = match open_resource_policy(&path) {
2514            Ok(file) => file,
2515            Err(error) if error.kind() == io::ErrorKind::NotFound => {
2516                return Ok(ResourcePolicyLoadEvidence {
2517                    policy: ResourcePolicy::everyday(),
2518                    source: ResourcePolicyLoadSource::MissingDefault,
2519                    warning: None,
2520                });
2521            }
2522            Err(error) => {
2523                if std::fs::symlink_metadata(&path)
2524                    .is_ok_and(|metadata| !metadata.is_file() || metadata.file_type().is_symlink())
2525                {
2526                    return Ok(corrupt_default(
2527                        "The saved resource policy could not be loaded because it is not a regular file.",
2528                    ));
2529                }
2530                return Err(error.into());
2531            }
2532        };
2533        let metadata = file.metadata()?;
2534        if !metadata.is_file() {
2535            return Ok(corrupt_default(
2536                "The saved resource policy could not be loaded because it is not a regular file.",
2537            ));
2538        }
2539        if metadata.len() > MAX_POLICY_BYTES {
2540            return Ok(corrupt_default(
2541                "The saved resource policy could not be loaded because it exceeds the size limit.",
2542            ));
2543        }
2544
2545        let mut raw = Vec::new();
2546        file.take(MAX_POLICY_BYTES + 1).read_to_end(&mut raw)?;
2547        if raw.len() as u64 > MAX_POLICY_BYTES {
2548            return Ok(corrupt_default(
2549                "The saved resource policy could not be loaded because it exceeds the size limit.",
2550            ));
2551        }
2552        let policy = match serde_json::from_slice::<ResourcePolicy>(&raw) {
2553            Ok(policy) => policy,
2554            Err(error) => {
2555                return Ok(corrupt_default(format!(
2556                    "The saved resource policy could not be loaded: {error}"
2557                )));
2558            }
2559        };
2560        if let Err(error) = policy.validate() {
2561            return Ok(corrupt_default(format!(
2562                "The saved resource policy could not be loaded: {error}"
2563            )));
2564        }
2565
2566        Ok(ResourcePolicyLoadEvidence {
2567            policy,
2568            source: ResourcePolicyLoadSource::Loaded,
2569            warning: None,
2570        })
2571    }
2572}
2573
2574fn open_resource_policy(path: &Path) -> io::Result<std::fs::File> {
2575    let mut options = std::fs::OpenOptions::new();
2576    options.read(true);
2577    #[cfg(unix)]
2578    {
2579        use std::os::unix::fs::OpenOptionsExt;
2580        options.custom_flags(libc::O_NOFOLLOW);
2581    }
2582    options.open(path)
2583}
2584
2585fn corrupt_default(warning: impl Into<String>) -> ResourcePolicyLoadEvidence {
2586    ResourcePolicyLoadEvidence {
2587        policy: ResourcePolicy::everyday(),
2588        source: ResourcePolicyLoadSource::CorruptDefault,
2589        warning: Some(warning.into()),
2590    }
2591}
2592
2593impl ResourcePolicyRepository for FileResourcePolicyRepository {
2594    fn load(&self) -> Result<ResourcePolicy, ResourcePolicyError> {
2595        Ok(self.load_with_evidence()?.policy)
2596    }
2597
2598    fn save(&self, policy: &ResourcePolicy) -> Result<(), ResourcePolicyError> {
2599        policy.validate()?;
2600        let _guard = mutation_lock()
2601            .lock()
2602            .unwrap_or_else(std::sync::PoisonError::into_inner);
2603        ensure_private_directory(&self.root)?;
2604
2605        let temp_path = unique_temp_path(&self.root);
2606        let result = (|| {
2607            let body = serde_json::to_vec_pretty(policy)?;
2608            let mut file = open_private_temp(&temp_path)?;
2609            file.write_all(&body)?;
2610            file.sync_all()?;
2611            atomic_replace(&temp_path, &self.path())?;
2612            car_secrets::harden_owner_only(&self.path());
2613            sync_directory(&self.root)?;
2614            Ok(())
2615        })();
2616        if result.is_err() {
2617            let _ = std::fs::remove_file(&temp_path);
2618        }
2619        result
2620    }
2621}
2622
2623fn mutation_lock() -> &'static Mutex<()> {
2624    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
2625    LOCK.get_or_init(|| Mutex::new(()))
2626}
2627
2628fn unique_temp_path(root: &Path) -> PathBuf {
2629    static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(0);
2630    let sequence = NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed);
2631    let epoch_nanos = std::time::SystemTime::now()
2632        .duration_since(std::time::UNIX_EPOCH)
2633        .map(|duration| duration.as_nanos())
2634        .unwrap_or(0);
2635    root.join(format!(
2636        ".{RESOURCE_POLICY_FILE}.{}.{}.{}.tmp",
2637        std::process::id(),
2638        epoch_nanos,
2639        sequence
2640    ))
2641}
2642
2643fn ensure_private_directory(path: &Path) -> io::Result<()> {
2644    std::fs::create_dir_all(path)?;
2645    #[cfg(unix)]
2646    {
2647        use std::os::unix::fs::PermissionsExt;
2648        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
2649    }
2650    car_secrets::harden_owner_only(path);
2651    Ok(())
2652}
2653
2654#[cfg(unix)]
2655fn open_private_temp(path: &Path) -> io::Result<std::fs::File> {
2656    use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
2657
2658    let file = std::fs::OpenOptions::new()
2659        .write(true)
2660        .create_new(true)
2661        .mode(0o600)
2662        .open(path)?;
2663    file.set_permissions(std::fs::Permissions::from_mode(0o600))?;
2664    Ok(file)
2665}
2666
2667#[cfg(not(unix))]
2668fn open_private_temp(path: &Path) -> io::Result<std::fs::File> {
2669    let file = std::fs::OpenOptions::new()
2670        .write(true)
2671        .create_new(true)
2672        .open(path)?;
2673    car_secrets::harden_owner_only(path);
2674    Ok(file)
2675}
2676
2677#[cfg(not(windows))]
2678fn atomic_replace(source: &Path, destination: &Path) -> io::Result<()> {
2679    std::fs::rename(source, destination)
2680}
2681
2682#[cfg(windows)]
2683fn atomic_replace(source: &Path, destination: &Path) -> io::Result<()> {
2684    use std::os::windows::ffi::OsStrExt;
2685
2686    const MOVEFILE_REPLACE_EXISTING: u32 = 0x1;
2687    const MOVEFILE_WRITE_THROUGH: u32 = 0x8;
2688    #[link(name = "kernel32")]
2689    unsafe extern "system" {
2690        fn MoveFileExW(existing: *const u16, new: *const u16, flags: u32) -> i32;
2691    }
2692    let source = source
2693        .as_os_str()
2694        .encode_wide()
2695        .chain(Some(0))
2696        .collect::<Vec<_>>();
2697    let destination = destination
2698        .as_os_str()
2699        .encode_wide()
2700        .chain(Some(0))
2701        .collect::<Vec<_>>();
2702    // SAFETY: both pointers address NUL-terminated UTF-16 buffers for the
2703    // duration of the call, and the flags request atomic replacement.
2704    let replaced = unsafe {
2705        MoveFileExW(
2706            source.as_ptr(),
2707            destination.as_ptr(),
2708            MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
2709        )
2710    };
2711    if replaced == 0 {
2712        Err(io::Error::last_os_error())
2713    } else {
2714        Ok(())
2715    }
2716}
2717
2718#[cfg(unix)]
2719fn sync_directory(path: &Path) -> io::Result<()> {
2720    std::fs::File::open(path)?.sync_all()
2721}
2722
2723#[cfg(not(unix))]
2724fn sync_directory(_path: &Path) -> io::Result<()> {
2725    Ok(())
2726}
2727
2728#[cfg(test)]
2729mod tests {
2730    use super::*;
2731
2732    fn hardware(
2733        total_ram_mb: u64,
2734        gpu_backend: crate::hardware::GpuBackend,
2735        vram_mb: Option<u64>,
2736    ) -> crate::hardware::HardwareInfo {
2737        crate::hardware::HardwareInfo {
2738            os: "test".into(),
2739            arch: "test".into(),
2740            cpu_cores: 8,
2741            total_ram_mb,
2742            gpu_backend,
2743            gpu_memory_mb: vram_mb,
2744            gpu_devices: Vec::new(),
2745            recommended_model: "fixture".into(),
2746            recommended_context: 4_096,
2747            max_model_mb: total_ram_mb,
2748        }
2749    }
2750
2751    #[test]
2752    fn profiles_compute_exact_32_gb_budgets() {
2753        let total = 32 * 1024;
2754
2755        assert_eq!(
2756            ResourcePolicy::everyday()
2757                .effective_budget(total)
2758                .configured_model_ceiling_mb,
2759            13_107
2760        );
2761        assert_eq!(
2762            ResourcePolicy::local_focused()
2763                .effective_budget(total)
2764                .configured_model_ceiling_mb,
2765            26_214
2766        );
2767        assert_eq!(
2768            ResourcePolicy::custom_gb(12.5)
2769                .unwrap()
2770                .effective_budget(total)
2771                .configured_model_ceiling_mb,
2772            12_800
2773        );
2774        assert_eq!(
2775            ResourcePolicy::everyday().recommendation_target_mb(total),
2776            6_553
2777        );
2778    }
2779
2780    #[test]
2781    fn model_memory_estimate_keeps_transient_margin_distinct_and_totals_exactly() {
2782        let catalog = crate::registry::builtin_catalog();
2783        let model = catalog
2784            .iter()
2785            .find(|model| model.id == "mlx/qwen3-4b:4bit")
2786            .unwrap();
2787        let estimate = estimate_model_memory(
2788            model,
2789            &hardware(32 * 1024, GpuBackend::Metal, None),
2790            RECOMMENDATION_CONTEXT_TOKENS,
2791        );
2792
2793        assert_eq!(estimate.evidence, ModelResourceEvidence::CatalogExact);
2794        assert_eq!(estimate.weights_mb, 2_400);
2795        assert_eq!(estimate.runtime_overhead_mb, 512);
2796        assert_eq!(estimate.transient_margin_mb, 1_024);
2797        assert_eq!(
2798            estimate.estimated_peak_mb,
2799            estimate.weights_mb
2800                + estimate.context_overhead_mb
2801                + estimate.runtime_overhead_mb
2802                + estimate.transient_margin_mb
2803        );
2804
2805        let measured = estimate_model_memory_with_measured_weights(
2806            model,
2807            &hardware(32 * 1024, GpuBackend::Metal, None),
2808            RECOMMENDATION_CONTEXT_TOKENS,
2809            Some(2_321),
2810        );
2811        assert_eq!(measured.weights_mb, 2_321);
2812        assert_eq!(measured.evidence, ModelResourceEvidence::FileSystemMeasured);
2813    }
2814
2815    fn qwen3_8b_model() -> ModelSchema {
2816        crate::registry::builtin_catalog()
2817            .into_iter()
2818            .find(|model| model.id == "mlx/qwen3-8b:4bit")
2819            .unwrap()
2820    }
2821
2822    fn qwen3_8b_geometry() -> KvGeometry {
2823        KvGeometry {
2824            num_hidden_layers: 36,
2825            num_key_value_heads: 8,
2826            head_dim: 128,
2827            kv_bytes_per_element: 2,
2828            hidden_size: Some(4_096),
2829        }
2830    }
2831
2832    #[test]
2833    fn kv_cache_qwen3_8b_geometry_matches_binary_mb_table() {
2834        let geometry = qwen3_8b_geometry();
2835        assert_eq!(kv_cache_mb_for(&geometry, 9_000), 1_266);
2836        assert_eq!(kv_cache_mb_for(&geometry, 13_096), 1_842);
2837        assert_eq!(prefill_activation_mb(Some(&geometry), 9_000), 282);
2838    }
2839
2840    #[test]
2841    fn kv_cache_unknown_geometry_uses_conservative_qwen3_8b_floor() {
2842        let model = qwen3_8b_model();
2843        let estimate =
2844            estimate_model_memory(&model, &hardware(32 * 1024, GpuBackend::Metal, None), 9_000);
2845
2846        assert_eq!(estimate.context_overhead_mb, 1_266);
2847    }
2848
2849    #[test]
2850    fn kv_cache_unknown_geometry_scales_models_larger_than_8b() {
2851        let mut model = qwen3_8b_model();
2852        model.param_count = "32B".into();
2853
2854        let estimate =
2855            estimate_model_memory(&model, &hardware(64 * 1024, GpuBackend::Metal, None), 9_000);
2856
2857        assert_eq!(estimate.context_overhead_mb, 5_063);
2858    }
2859
2860    #[test]
2861    fn kv_geometry_installed_config_drives_admission_context_charge() {
2862        let root = tempfile::tempdir().unwrap();
2863        let model = qwen3_8b_model();
2864        let model_dir = root.path().join(&model.name);
2865        std::fs::create_dir_all(&model_dir).unwrap();
2866        std::fs::write(
2867            model_dir.join("config.json"),
2868            r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32,"torch_dtype":"bfloat16"}"#,
2869        )
2870        .unwrap();
2871        let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2872            ResourcePolicy::local_focused(),
2873            hardware(32 * 1024, GpuBackend::Metal, None),
2874            Arc::new(FixedLiveMemoryProbe::known(24_000)),
2875            Some(root.path().to_path_buf()),
2876        );
2877
2878        let estimate = coordinator.preflight(&model, 9_000).estimate;
2879
2880        assert_eq!(estimate.context_overhead_mb, 1_548);
2881    }
2882
2883    /// The recommendation estimate takes the geometry admission reads, from
2884    /// the same file and through the same cache, so for an installed row the
2885    /// two publish one peak and only live memory separates their verdicts
2886    /// (car#2213).
2887    #[test]
2888    fn recommendation_estimate_uses_the_installed_geometry_admission_uses() {
2889        let root = tempfile::tempdir().unwrap();
2890        let model = qwen3_8b_model();
2891        let model_dir = root.path().join(&model.name);
2892        std::fs::create_dir_all(&model_dir).unwrap();
2893        std::fs::write(
2894            model_dir.join("config.json"),
2895            r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32,"torch_dtype":"bfloat16"}"#,
2896        )
2897        .unwrap();
2898        let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
2899        let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2900            ResourcePolicy::local_focused(),
2901            hardware.clone(),
2902            Arc::new(FixedLiveMemoryProbe::known(24_000)),
2903            Some(root.path().to_path_buf()),
2904        );
2905        let admitted = coordinator
2906            .preflight(&model, RECOMMENDATION_CONTEXT_TOKENS)
2907            .estimate
2908            .estimated_peak_mb;
2909        let peak = |geometry: &CheckpointGeometry| {
2910            estimate_model_memory_for_checkpoint(
2911                &model,
2912                &hardware,
2913                RECOMMENDATION_CONTEXT_TOKENS,
2914                geometry,
2915            )
2916            .estimated_peak_mb
2917        };
2918
2919        assert_eq!(
2920            peak(&CheckpointGeometry::read([&model], root.path())),
2921            admitted
2922        );
2923        assert_eq!(peak(&coordinator.checkpoint_geometry([&model])), admitted);
2924        // Not installed: the parameter-count fallback, which differs here.
2925        assert_ne!(peak(&CheckpointGeometry::default()), admitted);
2926        assert_eq!(
2927            peak(&CheckpointGeometry::default()),
2928            estimate_model_memory(&model, &hardware, RECOMMENDATION_CONTEXT_TOKENS)
2929                .estimated_peak_mb
2930        );
2931    }
2932
2933    #[test]
2934    fn kv_geometry_reads_a_multimodal_checkpoints_nested_text_config() {
2935        let root = tempfile::tempdir().unwrap();
2936        let model = qwen3_8b_model();
2937        let model_dir = root.path().join(&model.name);
2938        std::fs::create_dir_all(&model_dir).unwrap();
2939        // Same geometry as the flat test above, nested the way gemma-4 is.
2940        std::fs::write(
2941            model_dir.join("config.json"),
2942            r#"{"model_type":"gemma4_unified","torch_dtype":"bfloat16","vision_config":{},"text_config":{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32}}"#,
2943        )
2944        .unwrap();
2945        let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2946            ResourcePolicy::local_focused(),
2947            hardware(32 * 1024, GpuBackend::Metal, None),
2948            Arc::new(FixedLiveMemoryProbe::known(24_000)),
2949            Some(root.path().to_path_buf()),
2950        );
2951        assert_eq!(
2952            coordinator
2953                .preflight(&model, 9_000)
2954                .estimate
2955                .context_overhead_mb,
2956            1_548
2957        );
2958    }
2959
2960    #[test]
2961    fn kv_geometry_cache_forgets_missing_config_after_install() {
2962        let root = tempfile::tempdir().unwrap();
2963        let model = qwen3_8b_model();
2964        let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2965            ResourcePolicy::local_focused(),
2966            hardware(32 * 1024, GpuBackend::Metal, None),
2967            Arc::new(FixedLiveMemoryProbe::known(24_000)),
2968            Some(root.path().to_path_buf()),
2969        );
2970
2971        let before_install = coordinator.preflight(&model, 9_000).estimate;
2972        let model_dir = root.path().join(&model.name);
2973        std::fs::create_dir_all(&model_dir).unwrap();
2974        std::fs::write(
2975            model_dir.join("config.json"),
2976            r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32,"torch_dtype":"bfloat16"}"#,
2977        )
2978        .unwrap();
2979        let after_install = coordinator.preflight(&model, 9_000).estimate;
2980
2981        assert_eq!(before_install.context_overhead_mb, 1_266);
2982        assert_eq!(after_install.context_overhead_mb, 1_548);
2983    }
2984
2985    #[test]
2986    fn kv_geometry_derives_head_dim_and_memoizes_installed_config() {
2987        let root = tempfile::tempdir().unwrap();
2988        let model = qwen3_8b_model();
2989        let model_dir = root.path().join(&model.name);
2990        std::fs::create_dir_all(&model_dir).unwrap();
2991        let config_path = model_dir.join("config.json");
2992        std::fs::write(
2993            &config_path,
2994            r#"{"num_hidden_layers":36,"num_key_value_heads":8,"hidden_size":4096,"num_attention_heads":32}"#,
2995        )
2996        .unwrap();
2997        let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2998            ResourcePolicy::local_focused(),
2999            hardware(32 * 1024, GpuBackend::Metal, None),
3000            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3001            Some(root.path().to_path_buf()),
3002        );
3003
3004        let first = coordinator.preflight(&model, 9_000).estimate;
3005        // Same length and mtime: the stamp is unchanged, so the cached
3006        // geometry stands and the file is not parsed again.
3007        let modified = std::fs::metadata(&config_path).unwrap().modified().unwrap();
3008        let len = std::fs::metadata(&config_path).unwrap().len() as usize;
3009        std::fs::write(&config_path, "x".repeat(len)).unwrap();
3010        std::fs::File::options()
3011            .write(true)
3012            .open(&config_path)
3013            .unwrap()
3014            .set_modified(modified)
3015            .unwrap();
3016        let second = coordinator.preflight(&model, 9_000).estimate;
3017
3018        assert_eq!(first.context_overhead_mb, 1_548);
3019        assert_eq!(second.context_overhead_mb, first.context_overhead_mb);
3020    }
3021
3022    #[test]
3023    fn kv_geometry_cache_is_scoped_to_coordinator_models_dir() {
3024        let first_root = tempfile::tempdir().unwrap();
3025        let second_root = tempfile::tempdir().unwrap();
3026        let model = qwen3_8b_model();
3027        for (root, config) in [
3028            (
3029                first_root.path(),
3030                r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096}"#,
3031            ),
3032            (
3033                second_root.path(),
3034                r#"{"num_hidden_layers":2,"num_key_value_heads":2,"head_dim":64,"hidden_size":256}"#,
3035            ),
3036        ] {
3037            let model_dir = root.join(&model.name);
3038            std::fs::create_dir_all(&model_dir).unwrap();
3039            std::fs::write(model_dir.join("config.json"), config).unwrap();
3040        }
3041        let coordinator = |models_dir: &Path| {
3042            LocalAdmissionCoordinator::with_probe_and_models_dir(
3043                ResourcePolicy::local_focused(),
3044                hardware(32 * 1024, GpuBackend::Metal, None),
3045                Arc::new(FixedLiveMemoryProbe::known(24_000)),
3046                Some(models_dir.to_path_buf()),
3047            )
3048        };
3049
3050        // Both coordinators look up the same model id. Their distinct results
3051        // prove geometry memoization belongs to the coordinator/models root,
3052        // rather than leaking one test or engine's config across the process.
3053        assert_eq!(
3054            coordinator(first_root.path())
3055                .preflight(&model, 9_000)
3056                .estimate
3057                .context_overhead_mb,
3058            1_548
3059        );
3060        assert_eq!(
3061            coordinator(second_root.path())
3062                .preflight(&model, 9_000)
3063                .estimate
3064                .context_overhead_mb,
3065            27
3066        );
3067    }
3068
3069    #[test]
3070    fn kv_geometry_malformed_config_falls_back_without_panicking() {
3071        let root = tempfile::tempdir().unwrap();
3072        let model = qwen3_8b_model();
3073        let model_dir = root.path().join(&model.name);
3074        std::fs::create_dir_all(&model_dir).unwrap();
3075        std::fs::write(model_dir.join("config.json"), b"not json").unwrap();
3076        let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
3077            ResourcePolicy::local_focused(),
3078            hardware(32 * 1024, GpuBackend::Metal, None),
3079            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3080            Some(root.path().to_path_buf()),
3081        );
3082
3083        let estimate = coordinator.preflight(&model, 9_000).estimate;
3084
3085        assert_eq!(estimate.context_overhead_mb, 1_266);
3086    }
3087
3088    #[test]
3089    fn kv_cache_estimate_is_monotone_and_keeps_non_context_terms_fixed() {
3090        let model = qwen3_8b_model();
3091        let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
3092        let short = estimate_model_memory(&model, &hardware, 9_000);
3093        let long = estimate_model_memory(&model, &hardware, 13_096);
3094
3095        assert!(long.context_overhead_mb >= short.context_overhead_mb);
3096        assert_eq!(short.weights_mb, long.weights_mb);
3097        assert_eq!(short.runtime_overhead_mb, long.runtime_overhead_mb);
3098        assert_eq!(short.transient_margin_mb, long.transient_margin_mb);
3099    }
3100
3101    #[test]
3102    fn custom_zero_disables_loads_and_overlarge_values_clamp_below_emergency_reserve() {
3103        let total = 32 * 1024;
3104
3105        assert_eq!(
3106            ResourcePolicy::custom_gb(0.0)
3107                .unwrap()
3108                .effective_budget(total)
3109                .effective_new_load_ceiling_mb,
3110            0
3111        );
3112        let result = ResourcePolicy::custom_gb(99.0)
3113            .unwrap()
3114            .effective_budget(total);
3115        assert_eq!(result.emergency_reserve_mb, 3_276);
3116        assert_eq!(result.configured_model_ceiling_mb, total - 3_276);
3117        assert!(result.normalization_notice.is_some());
3118    }
3119
3120    #[test]
3121    fn custom_gigabytes_reject_non_half_steps_negative_and_non_finite_values() {
3122        for invalid in [10.3, -0.5, f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
3123            assert!(
3124                ResourcePolicy::custom_gb(invalid).is_err(),
3125                "accepted {invalid:?}"
3126            );
3127        }
3128
3129        assert_eq!(
3130            ResourcePolicy::custom_gb(10.5).unwrap().custom_max_model_mb,
3131            Some(10_752)
3132        );
3133    }
3134
3135    #[test]
3136    fn repository_round_trips_exact_half_gb_and_uses_private_atomic_files() {
3137        let dir = tempfile::tempdir().unwrap();
3138        let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3139
3140        repository
3141            .save(&ResourcePolicy::custom_gb(10.5).unwrap())
3142            .unwrap();
3143
3144        assert_eq!(
3145            repository.path(),
3146            dir.path().join("model-resource-policy.json")
3147        );
3148        assert_eq!(repository.load().unwrap().custom_max_model_mb, Some(10_752));
3149        assert_private_mode(&repository.path(), 0o600);
3150        assert_private_mode(dir.path(), 0o700);
3151        assert_no_temp_files(dir.path());
3152    }
3153
3154    #[test]
3155    fn repository_missing_or_corrupt_file_falls_back_without_deleting_source() {
3156        let dir = tempfile::tempdir().unwrap();
3157        let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3158        assert_eq!(repository.load().unwrap(), ResourcePolicy::everyday());
3159
3160        let corrupt = br#"{"profile":"custom","custom_max_model_mb":"broken"}"#;
3161        std::fs::write(repository.path(), corrupt).unwrap();
3162
3163        assert_eq!(repository.load().unwrap(), ResourcePolicy::everyday());
3164        assert_eq!(std::fs::read(repository.path()).unwrap(), corrupt);
3165    }
3166
3167    #[test]
3168    fn repository_rejects_invalid_policy_shapes_before_writing() {
3169        let dir = tempfile::tempdir().unwrap();
3170        let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3171        let invalid = [
3172            ResourcePolicy {
3173                profile: ResourceProfile::Custom,
3174                custom_max_model_mb: None,
3175            },
3176            ResourcePolicy {
3177                profile: ResourceProfile::Custom,
3178                custom_max_model_mb: Some(1),
3179            },
3180            ResourcePolicy {
3181                profile: ResourceProfile::Custom,
3182                custom_max_model_mb: Some(513),
3183            },
3184            ResourcePolicy {
3185                profile: ResourceProfile::Everyday,
3186                custom_max_model_mb: Some(512),
3187            },
3188            ResourcePolicy {
3189                profile: ResourceProfile::LocalFocused,
3190                custom_max_model_mb: Some(512),
3191            },
3192        ];
3193
3194        for policy in invalid {
3195            let error = repository.save(&policy).unwrap_err();
3196            assert!(matches!(error, ResourcePolicyError::InvalidPolicy { .. }));
3197            assert!(!repository.path().exists());
3198        }
3199    }
3200
3201    #[test]
3202    fn load_evidence_distinguishes_loaded_missing_and_corrupt_defaults() {
3203        let dir = tempfile::tempdir().unwrap();
3204        let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3205
3206        let missing = repository.load_with_evidence().unwrap();
3207        assert_eq!(missing.policy, ResourcePolicy::everyday());
3208        assert_eq!(missing.source, ResourcePolicyLoadSource::MissingDefault);
3209        assert!(missing.warning.is_none());
3210
3211        repository.save(&ResourcePolicy::local_focused()).unwrap();
3212        let loaded = repository.load_with_evidence().unwrap();
3213        assert_eq!(loaded.policy, ResourcePolicy::local_focused());
3214        assert_eq!(loaded.source, ResourcePolicyLoadSource::Loaded);
3215        assert!(loaded.warning.is_none());
3216
3217        let corrupt = br#"{"profile":"custom","custom_max_model_mb":"broken"}"#;
3218        std::fs::write(repository.path(), corrupt).unwrap();
3219        let recovered = repository.load_with_evidence().unwrap();
3220        assert_eq!(recovered.policy, ResourcePolicy::everyday());
3221        assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
3222        assert!(recovered
3223            .warning
3224            .as_deref()
3225            .is_some_and(|warning| { warning.contains("could not be loaded") }));
3226        assert_eq!(std::fs::read(repository.path()).unwrap(), corrupt);
3227    }
3228
3229    #[test]
3230    fn corrupt_evidence_covers_unknown_fields_invalid_shapes_and_oversized_files() {
3231        let dir = tempfile::tempdir().unwrap();
3232        let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3233        let corrupt_documents = [
3234            br#"{"#.to_vec(),
3235            br#"{"profile":"everyday","custom_max_model_mb":null,"extra":true}"#.to_vec(),
3236            br#"{"profile":"custom","custom_max_model_mb":null}"#.to_vec(),
3237            br#"{"profile":"custom","custom_max_model_mb":1}"#.to_vec(),
3238            br#"{"profile":"custom","custom_max_model_mb":513}"#.to_vec(),
3239            br#"{"profile":"everyday","custom_max_model_mb":512}"#.to_vec(),
3240            br#"{"profile":"local_focused","custom_max_model_mb":512}"#.to_vec(),
3241            vec![b' '; MAX_POLICY_BYTES as usize + 1],
3242        ];
3243
3244        for document in corrupt_documents {
3245            std::fs::write(repository.path(), &document).unwrap();
3246            let recovered = repository.load_with_evidence().unwrap();
3247            assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
3248            assert!(recovered.warning.is_some());
3249            assert_eq!(std::fs::read(repository.path()).unwrap(), document);
3250        }
3251    }
3252
3253    #[cfg(unix)]
3254    #[test]
3255    fn non_regular_policy_source_is_reported_as_corrupt_without_following_it() {
3256        use std::os::unix::fs::symlink;
3257
3258        let dir = tempfile::tempdir().unwrap();
3259        let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3260        let target = dir.path().join("target.json");
3261        std::fs::write(
3262            &target,
3263            br#"{"profile":"local_focused","custom_max_model_mb":null}"#,
3264        )
3265        .unwrap();
3266        symlink(&target, repository.path()).unwrap();
3267
3268        let recovered = repository.load_with_evidence().unwrap();
3269        assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
3270        assert!(recovered.warning.is_some());
3271        assert!(repository.path().is_symlink());
3272    }
3273
3274    #[test]
3275    fn concurrent_saves_leave_one_complete_document_and_no_staging_files() {
3276        let dir = tempfile::tempdir().unwrap();
3277        let repository =
3278            std::sync::Arc::new(FileResourcePolicyRepository::new(dir.path().to_path_buf()));
3279        let mut writers = Vec::new();
3280        for index in 0..16_u64 {
3281            let repository = repository.clone();
3282            writers.push(std::thread::spawn(move || {
3283                repository
3284                    .save(&ResourcePolicy {
3285                        profile: ResourceProfile::Custom,
3286                        custom_max_model_mb: Some(index * 512),
3287                    })
3288                    .unwrap();
3289            }));
3290        }
3291        for writer in writers {
3292            writer.join().unwrap();
3293        }
3294
3295        let loaded = repository.load().unwrap();
3296        assert_eq!(loaded.profile, ResourceProfile::Custom);
3297        assert!(loaded.custom_max_model_mb.unwrap().is_multiple_of(512));
3298        assert_no_temp_files(dir.path());
3299    }
3300
3301    #[test]
3302    fn cuda_uses_separate_vram_fit_and_host_ram_policy() {
3303        let hardware = hardware(
3304            64 * 1024,
3305            crate::hardware::GpuBackend::Cuda,
3306            Some(12 * 1024),
3307        );
3308
3309        let evidence = evaluate_resources(&hardware, &ResourcePolicy::everyday());
3310
3311        assert_eq!(evidence.host_memory.configured_model_ceiling_mb, 26_214);
3312        let accelerator = evidence.accelerator_memory.unwrap();
3313        assert_eq!(accelerator.total_mb, 12 * 1024);
3314        assert_eq!(accelerator.budget_mb, 12 * 1024);
3315        assert_ne!(
3316            accelerator.budget_mb,
3317            evidence.host_memory.configured_model_ceiling_mb
3318        );
3319    }
3320
3321    #[test]
3322    fn resource_preflight_zero_custom_budget_blocks_new_load() {
3323        let coordinator = LocalAdmissionCoordinator::with_probe(
3324            ResourcePolicy::custom_gb(0.0).unwrap(),
3325            hardware(32 * 1024, GpuBackend::Metal, None),
3326            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3327        );
3328        let model = crate::registry::builtin_catalog()
3329            .into_iter()
3330            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3331            .unwrap();
3332
3333        let preflight = coordinator.preflight(&model, 2_048);
3334        assert_eq!(preflight.verdict, LocalLoadVerdict::DisabledByPolicy);
3335    }
3336
3337    #[test]
3338    fn resource_preflight_zero_budget_does_not_kill_resident_inference() {
3339        let coordinator = LocalAdmissionCoordinator::with_probe(
3340            ResourcePolicy::custom_gb(0.0).unwrap(),
3341            hardware(32 * 1024, GpuBackend::Metal, None),
3342            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3343        );
3344        let model = crate::registry::builtin_catalog()
3345            .into_iter()
3346            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3347            .unwrap();
3348        coordinator.mark_resident(&model.id, 2_400);
3349
3350        let preflight = coordinator.preflight(&model, 2_048);
3351        assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
3352        assert_eq!(
3353            preflight.estimated_incremental_mb,
3354            preflight.estimate.context_overhead_mb
3355                + preflight.estimate.runtime_overhead_mb
3356                + preflight.estimate.transient_margin_mb
3357        );
3358    }
3359
3360    #[test]
3361    fn resource_preflight_resident_weights_are_incremental_only() {
3362        let coordinator = LocalAdmissionCoordinator::with_probe(
3363            ResourcePolicy::custom_gb(8.0).unwrap(),
3364            hardware(32 * 1024, GpuBackend::Metal, None),
3365            Arc::new(FixedLiveMemoryProbe::known(6_000)),
3366        );
3367        let model = crate::registry::builtin_catalog()
3368            .into_iter()
3369            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3370            .unwrap();
3371        coordinator.mark_resident(&model.id, 2_400);
3372
3373        let preflight = coordinator.preflight(&model, 2_048);
3374        assert_eq!(
3375            preflight.estimated_incremental_mb,
3376            preflight.estimate.context_overhead_mb
3377                + preflight.estimate.runtime_overhead_mb
3378                + preflight.estimate.transient_margin_mb
3379        );
3380        assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
3381    }
3382
3383    #[test]
3384    fn resource_preflight_unavailable_live_probe_is_explicit() {
3385        let coordinator = LocalAdmissionCoordinator::with_probe(
3386            ResourcePolicy::everyday(),
3387            hardware(32 * 1024, GpuBackend::Metal, None),
3388            Arc::new(FixedLiveMemoryProbe::unknown()),
3389        );
3390        let model = crate::registry::builtin_catalog()
3391            .into_iter()
3392            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3393            .unwrap();
3394
3395        let preflight = coordinator.preflight(&model, 2_048);
3396        assert_eq!(preflight.live_available_mb, None);
3397        assert_eq!(preflight.verdict, LocalLoadVerdict::LiveMemoryUnknown);
3398    }
3399
3400    #[test]
3401    fn resource_preflight_simultaneous_reservations_are_atomic() {
3402        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3403            ResourcePolicy::custom_gb(5.0).unwrap(),
3404            hardware(32 * 1024, GpuBackend::Metal, None),
3405            Arc::new(FixedLiveMemoryProbe::known(6_000)),
3406        ));
3407        let model = crate::registry::builtin_catalog()
3408            .into_iter()
3409            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3410            .unwrap();
3411        coordinator.mark_resident(&model.id, 2_400);
3412
3413        let first = coordinator.reserve(&model, 2_048).unwrap();
3414        let second = coordinator.reserve(&model, 2_048).unwrap_err();
3415        assert_eq!(
3416            second.preflight.verdict,
3417            LocalLoadVerdict::ExceedsConfiguredCeiling
3418        );
3419        drop(first);
3420        assert!(coordinator.reserve(&model, 2_048).is_ok());
3421    }
3422
3423    #[test]
3424    fn distinct_state_roots_reserve_atomically_against_one_machine_ledger() {
3425        let machine_ledger = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
3426        let policy = ResourcePolicy::custom_gb(6.0).unwrap();
3427        let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
3428        let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3429            policy.clone(),
3430            hardware.clone(),
3431            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3432            machine_ledger.clone(),
3433        ));
3434        let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3435            policy,
3436            hardware,
3437            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3438            machine_ledger,
3439        ));
3440        let mut model = crate::registry::builtin_catalog()
3441            .into_iter()
3442            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3443            .unwrap();
3444        model.cost.ram_mb = Some(2 * 1024);
3445        model.cost.size_mb = Some(2 * 1024);
3446        let barrier = Arc::new(std::sync::Barrier::new(2));
3447
3448        let attempts = [first, second].map(|coordinator| {
3449            let model = model.clone();
3450            let barrier = barrier.clone();
3451            std::thread::spawn(move || {
3452                barrier.wait();
3453                coordinator.reserve(&model, 0)
3454            })
3455        });
3456        let outcomes = attempts.map(|attempt| attempt.join().unwrap());
3457
3458        assert_eq!(outcomes.iter().filter(|result| result.is_ok()).count(), 1);
3459        assert_eq!(
3460            outcomes
3461                .iter()
3462                .find_map(|result| result.as_ref().err())
3463                .expect("one cross-root request must be blocked")
3464                .preflight
3465                .verdict,
3466            LocalLoadVerdict::ExceedsConfiguredCeiling
3467        );
3468    }
3469
3470    #[test]
3471    fn resource_preflight_cuda_charges_weights_to_vram_and_overhead_to_host() {
3472        let coordinator = LocalAdmissionCoordinator::with_probe(
3473            ResourcePolicy::custom_gb(2.0).unwrap(),
3474            hardware(64 * 1024, GpuBackend::Cuda, Some(12 * 1024)),
3475            Arc::new(FixedLiveMemoryProbe::known(20_000)),
3476        );
3477        let model = crate::registry::builtin_catalog()
3478            .into_iter()
3479            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3480            .unwrap();
3481
3482        let preflight = coordinator.preflight(&model, 2_048);
3483        assert_eq!(preflight.resident_model_mb, 0);
3484        assert_eq!(preflight.accelerator_total_mb, Some(12 * 1024));
3485        assert_eq!(
3486            preflight.accelerator_incremental_mb,
3487            Some(preflight.estimate.weights_mb)
3488        );
3489        assert_eq!(
3490            preflight.estimated_incremental_mb,
3491            preflight.estimate.context_overhead_mb
3492                + preflight.estimate.runtime_overhead_mb
3493                + preflight.estimate.transient_margin_mb
3494        );
3495        assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
3496    }
3497
3498    #[test]
3499    fn resource_preflight_cuda_vram_reservations_are_atomic() {
3500        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3501            ResourcePolicy::custom_gb(8.0).unwrap(),
3502            hardware(64 * 1024, GpuBackend::Cuda, Some(3_000)),
3503            Arc::new(FixedLiveMemoryProbe::known(20_000)),
3504        ));
3505        let model = crate::registry::builtin_catalog()
3506            .into_iter()
3507            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3508            .unwrap();
3509
3510        let first = coordinator.reserve(&model, 2_048).unwrap();
3511        assert_eq!(
3512            coordinator
3513                .reserve(&model, 2_048)
3514                .unwrap_err()
3515                .preflight
3516                .verdict,
3517            LocalLoadVerdict::ExceedsConfiguredCeiling
3518        );
3519        drop(first);
3520        assert!(coordinator.reserve(&model, 2_048).is_ok());
3521    }
3522
3523    #[test]
3524    fn resource_preflight_model_maintenance_races_atomically_with_reserve() {
3525        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3526            ResourcePolicy::everyday(),
3527            hardware(32 * 1024, GpuBackend::Metal, None),
3528            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3529        ));
3530        let model = crate::registry::builtin_catalog()
3531            .into_iter()
3532            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3533            .unwrap();
3534
3535        let active = coordinator.reserve(&model, 2_048).unwrap();
3536        assert!(matches!(
3537            coordinator.begin_model_maintenance(&model.id),
3538            Err(ModelMaintenanceError::ModelInUse(_))
3539        ));
3540        drop(active);
3541
3542        let maintenance = coordinator.begin_model_maintenance(&model.id).unwrap();
3543        assert_eq!(
3544            coordinator
3545                .reserve(&model, 2_048)
3546                .unwrap_err()
3547                .preflight
3548                .verdict,
3549            LocalLoadVerdict::ModelMaintenance
3550        );
3551        drop(maintenance);
3552        assert!(coordinator.reserve(&model, 2_048).is_ok());
3553    }
3554
3555    #[test]
3556    fn measured_weights_are_rechecked_atomically_before_allocation() {
3557        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3558            ResourcePolicy::custom_gb(5.0).unwrap(),
3559            hardware(32 * 1024, GpuBackend::Metal, None),
3560            Arc::new(FixedLiveMemoryProbe::known(16_000)),
3561        ));
3562        let model = crate::registry::builtin_catalog()
3563            .into_iter()
3564            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3565            .unwrap();
3566        let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
3567
3568        let blocked = reservation
3569            .reconcile_measured_weights(6 * 1024 * 1024 * 1024)
3570            .unwrap_err();
3571
3572        assert_eq!(
3573            blocked.preflight.verdict,
3574            LocalLoadVerdict::ExceedsConfiguredCeiling
3575        );
3576        assert_eq!(
3577            coordinator.preflight(&model, 2_048).active_reservations_mb,
3578            reservation.reserved_incremental_mb(),
3579            "a rejected resize must not mutate the live reservation"
3580        );
3581    }
3582
3583    #[test]
3584    fn cache_publication_transfers_cold_reservation_without_double_counting() {
3585        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3586            ResourcePolicy::custom_gb(8.0).unwrap(),
3587            hardware(32 * 1024, GpuBackend::Metal, None),
3588            Arc::new(FixedLiveMemoryProbe::known(16_000)),
3589        ));
3590        let model = crate::registry::builtin_catalog()
3591            .into_iter()
3592            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3593            .unwrap();
3594        let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
3595        let measured_bytes = 3 * 1024 * 1024 * 1024_u64;
3596
3597        reservation
3598            .reconcile_measured_weights(measured_bytes)
3599            .unwrap();
3600        reservation.publish_resident_weights(measured_bytes);
3601        let after = coordinator.preflight(&model, 2_048);
3602        let request_overhead = after.estimate.context_overhead_mb
3603            + after.estimate.runtime_overhead_mb
3604            + after.estimate.transient_margin_mb;
3605
3606        assert_eq!(after.resident_model_mb, 3 * 1024);
3607        assert_eq!(after.active_reservations_mb, request_overhead);
3608        assert_eq!(reservation.reserved_incremental_mb(), request_overhead);
3609    }
3610
3611    #[test]
3612    fn resident_publication_never_reprobes_or_rejects_after_allocation() {
3613        struct CountingProbe {
3614            calls: AtomicU64,
3615            available_mb: u64,
3616        }
3617
3618        impl LiveMemoryProbe for CountingProbe {
3619            fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
3620                self.calls.fetch_add(1, Ordering::Relaxed);
3621                Ok(Some(self.available_mb))
3622            }
3623        }
3624
3625        let probe = Arc::new(CountingProbe {
3626            calls: AtomicU64::new(0),
3627            available_mb: 24_000,
3628        });
3629        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3630            ResourcePolicy::custom_gb(8.0).unwrap(),
3631            hardware(32 * 1024, GpuBackend::Metal, None),
3632            probe.clone(),
3633        ));
3634        let model = crate::registry::builtin_catalog()
3635            .into_iter()
3636            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3637            .unwrap();
3638        let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
3639        let measured = 3 * 1024 * 1024 * 1024_u64;
3640        reservation.reconcile_measured_weights(measured).unwrap();
3641        let calls_before_publish = probe.calls.load(Ordering::Relaxed);
3642
3643        reservation.publish_resident_weights(measured + 512 * 1024 * 1024);
3644
3645        assert_eq!(probe.calls.load(Ordering::Relaxed), calls_before_publish);
3646        assert!(coordinator.is_resident(&model.id));
3647        assert_eq!(
3648            coordinator
3649                .state
3650                .lock()
3651                .unwrap_or_else(std::sync::PoisonError::into_inner)
3652                .resident_models
3653                .get(&model.id)
3654                .map(|allocation| allocation.weights_mb),
3655            Some(3 * 1024 + 512)
3656        );
3657    }
3658
3659    #[test]
3660    fn simultaneous_cold_reservations_converge_on_one_resident_allocation() {
3661        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3662            ResourcePolicy::custom_gb(16.0).unwrap(),
3663            hardware(32 * 1024, GpuBackend::Metal, None),
3664            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3665        ));
3666        let model = crate::registry::builtin_catalog()
3667            .into_iter()
3668            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3669            .unwrap();
3670        let mut first = coordinator.reserve(&model, 2_048).unwrap();
3671        let mut second = coordinator.reserve(&model, 2_048).unwrap();
3672        let measured = 3 * 1024 * 1024 * 1024_u64;
3673
3674        first.reconcile_measured_weights(measured).unwrap();
3675        second.reconcile_measured_weights(measured).unwrap();
3676        first.publish_resident_weights(measured);
3677        second.publish_resident_weights(measured);
3678
3679        let after = coordinator.preflight(&model, 2_048);
3680        let per_request_overhead = after.estimate.context_overhead_mb
3681            + after.estimate.runtime_overhead_mb
3682            + after.estimate.transient_margin_mb;
3683        assert_eq!(after.resident_model_mb, 3 * 1024);
3684        assert_eq!(
3685            after.active_reservations_mb,
3686            per_request_overhead * 2,
3687            "peer publication must release only this reservation's redundant cold weights"
3688        );
3689    }
3690
3691    #[test]
3692    fn resident_eviction_before_allocation_repromotes_reservation_to_cold() {
3693        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3694            ResourcePolicy::custom_gb(8.0).unwrap(),
3695            hardware(32 * 1024, GpuBackend::Cpu, None),
3696            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3697        ));
3698        let model = crate::registry::builtin_catalog()
3699            .into_iter()
3700            .find(|model| model.id == "mlx/qwen3-4b:4bit")
3701            .unwrap();
3702        let measured = 3 * 1024 * 1024 * 1024_u64;
3703        let mut first = coordinator.reserve(&model, 2_048).unwrap();
3704        first.reconcile_measured_weights(measured).unwrap();
3705        first.publish_resident_weights(measured);
3706        drop(first);
3707        let mut replacement = coordinator.reserve(&model, 2_048).unwrap();
3708        coordinator.mark_evicted(&model.id);
3709
3710        let preflight = replacement.reconcile_measured_weights(measured).unwrap();
3711        assert!(preflight.estimated_incremental_mb >= 3 * 1024);
3712    }
3713
3714    #[test]
3715    fn maintenance_catalog_alias_blocks_provider_alias_reservation() {
3716        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3717            ResourcePolicy::custom_gb(8.0).unwrap(),
3718            hardware(32 * 1024, GpuBackend::Cpu, None),
3719            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3720        ));
3721        coordinator
3722            .register_model_aliases("mlx/kokoro-82m:6bit", ["mlx-community/Kokoro-82M-6bit"]);
3723        let _maintenance = coordinator
3724            .begin_model_maintenance("mlx/kokoro-82m:6bit")
3725            .unwrap();
3726        let mut provider_schema = crate::registry::builtin_catalog()
3727            .into_iter()
3728            .find(|model| model.is_local())
3729            .expect("one local schema");
3730        provider_schema.id = "mlx-community/Kokoro-82M-6bit".into();
3731
3732        let blocked = coordinator.reserve(&provider_schema, 512).unwrap_err();
3733        assert_eq!(
3734            blocked.preflight.verdict,
3735            LocalLoadVerdict::ModelMaintenance
3736        );
3737        assert_eq!(blocked.preflight.model_id, "mlx/kokoro-82m:6bit");
3738    }
3739
3740    #[test]
3741    fn late_alias_registration_rekeys_active_and_resident_state_atomically() {
3742        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3743            ResourcePolicy::custom_gb(8.0).unwrap(),
3744            hardware(32 * 1024, GpuBackend::Cpu, None),
3745            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3746        ));
3747        let mut schema = crate::registry::builtin_catalog()
3748            .into_iter()
3749            .find(|model| model.is_local())
3750            .unwrap();
3751        schema.id = "provider/artifact".into();
3752        let active = coordinator.reserve(&schema, 0).unwrap();
3753
3754        coordinator.register_model_aliases("catalog/model:default", ["provider/artifact"]);
3755        assert!(matches!(
3756            coordinator.begin_model_maintenance("catalog/model:default"),
3757            Err(ModelMaintenanceError::ModelInUse(_))
3758        ));
3759        drop(active);
3760
3761        let maintenance = coordinator
3762            .begin_model_maintenance("catalog/model:default")
3763            .unwrap();
3764        let error = coordinator.reserve(&schema, 0).unwrap_err();
3765        assert_eq!(error.preflight.verdict, LocalLoadVerdict::ModelMaintenance);
3766        drop(maintenance);
3767    }
3768
3769    #[test]
3770    fn process_allocations_teardown_by_exact_owner_without_sibling_erasure() {
3771        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3772            ResourcePolicy::custom_gb(16.0).unwrap(),
3773            hardware(32 * 1024, GpuBackend::Cpu, None),
3774            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3775        ));
3776        let logical = "local/shared-model";
3777        let worker = worker_process_allocation_id(logical);
3778        let vllm = vllm_process_allocation_id(logical);
3779        let bytes = 512 * 1024 * 1024_u64;
3780        let mut worker_load = coordinator
3781            .reserve_measured_host_allocation(logical, &worker, bytes, 0)
3782            .unwrap();
3783        worker_load.publish_resident_weights_as(&worker, bytes);
3784        drop(worker_load);
3785        let mut vllm_load = coordinator
3786            .reserve_measured_host_allocation(logical, &vllm, bytes, 0)
3787            .unwrap();
3788        vllm_load.publish_resident_weights_as(&vllm, bytes);
3789        drop(vllm_load);
3790        assert_eq!(coordinator.resident_model_mb(), 1024);
3791
3792        coordinator.mark_teardown_pending_allocation(logical, &worker);
3793        coordinator.mark_teardown_pending_allocation(logical, &vllm);
3794        coordinator.finish_teardown_allocation(logical, &worker);
3795        assert!(coordinator.teardown_pending(logical));
3796        assert_eq!(
3797            coordinator.resident_allocation_ids(logical),
3798            vec![vllm.clone()]
3799        );
3800        assert_eq!(coordinator.resident_model_mb(), 512);
3801
3802        coordinator.finish_teardown_allocation(logical, &vllm);
3803        assert!(!coordinator.teardown_pending(logical));
3804        assert!(!coordinator.is_resident(logical));
3805    }
3806
3807    #[test]
3808    fn replacement_generation_reconcile_removes_peer_resident_discount_before_load() {
3809        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3810            ResourcePolicy::custom_gb(1.0).unwrap(),
3811            hardware(32 * 1024, GpuBackend::Cpu, None),
3812            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3813        ));
3814        let logical = "local/replaced-worker";
3815        let bytes = 512 * 1024 * 1024_u64;
3816        let mut old = coordinator
3817            .reserve_measured_host_allocation(logical, "worker:old", bytes, 0)
3818            .unwrap();
3819        old.publish_resident_weights_as("worker:old", bytes);
3820        drop(old);
3821
3822        let mut replacement = coordinator
3823            .reserve_measured_host(logical, bytes, 0)
3824            .expect("logical preflight initially sees the old resident");
3825        assert_eq!(replacement.reserved_incremental_mb(), 0);
3826        replacement.bind_allocation_id("worker:new");
3827        replacement
3828            .reconcile_measured_weights(bytes)
3829            .expect("the exact replacement generation fits by itself");
3830        assert_eq!(replacement.reserved_incremental_mb(), 512);
3831        assert!(
3832            coordinator
3833                .reserve_measured_host("different/model", bytes, 0)
3834                .is_err(),
3835            "a second cold model must see both old residency and the replacement generation"
3836        );
3837
3838        drop(replacement);
3839        assert!(coordinator
3840            .reserve_measured_host("different/model", bytes, 0)
3841            .is_ok());
3842    }
3843
3844    #[test]
3845    fn pre_ack_pending_charge_blocks_other_scope_until_exact_exit_ack() {
3846        let machine = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
3847        let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3848            ResourcePolicy::custom_gb(2.0).unwrap(),
3849            hardware(32 * 1024, GpuBackend::Cpu, None),
3850            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3851            machine.clone(),
3852        ));
3853        let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3854            ResourcePolicy::custom_gb(2.0).unwrap(),
3855            hardware(32 * 1024, GpuBackend::Cpu, None),
3856            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3857            machine,
3858        ));
3859        let logical = "managed/starting";
3860        let allocation = worker_process_allocation_id(logical);
3861        let cold = first
3862            .reserve_measured_host_allocation(logical, &allocation, 1024 * 1024 * 1024, 0)
3863            .unwrap();
3864        first.mark_teardown_pending_allocation_with_charge(
3865            logical,
3866            &allocation,
3867            cold.reconciled_weights_bytes(),
3868        );
3869        drop(cold);
3870
3871        let blocked = second
3872            .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3873            .unwrap_err();
3874        assert_eq!(
3875            blocked.preflight.verdict,
3876            LocalLoadVerdict::ExceedsConfiguredCeiling
3877        );
3878        first.finish_teardown_allocation(logical, "worker:unrelated-sibling");
3879        assert!(first.teardown_pending(logical));
3880        assert_eq!(first.resident_model_mb(), 1024);
3881
3882        first.finish_teardown_allocation(logical, &allocation);
3883        assert!(!first.teardown_pending(logical));
3884        assert!(second
3885            .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3886            .is_ok());
3887    }
3888
3889    #[test]
3890    fn detached_native_lease_keeps_machine_charge_after_request_cancellation() {
3891        let machine = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
3892        let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3893            ResourcePolicy::custom_gb(2.0).unwrap(),
3894            hardware(32 * 1024, GpuBackend::Cpu, None),
3895            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3896            machine.clone(),
3897        ));
3898        let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3899            ResourcePolicy::custom_gb(2.0).unwrap(),
3900            hardware(32 * 1024, GpuBackend::Cpu, None),
3901            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3902            machine,
3903        ));
3904        let reservation = first
3905            .reserve_measured_host("detached/model", 1024 * 1024 * 1024, 0)
3906            .unwrap();
3907        let detached = reservation.detached_lease();
3908        drop(reservation);
3909        assert!(second
3910            .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3911            .is_err());
3912        drop(detached);
3913        assert!(second
3914            .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3915            .is_ok());
3916    }
3917
3918    #[test]
3919    fn normal_awaited_detached_work_shares_one_charge_with_its_request() {
3920        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3921            ResourcePolicy::custom_gb(1.0).unwrap(),
3922            hardware(32 * 1024, GpuBackend::Cpu, None),
3923            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3924        ));
3925        let request = coordinator
3926            .reserve_measured_host("native/model-a", 512 * 1024 * 1024, 0)
3927            .unwrap();
3928        let detached = request.detached_lease();
3929
3930        let peer = coordinator
3931            .reserve_measured_host("native/model-b", 512 * 1024 * 1024, 0)
3932            .expect("request + its detached job are one 512 MB allocation, not two");
3933        assert_eq!(
3934            coordinator
3935                .state
3936                .lock()
3937                .unwrap_or_else(std::sync::PoisonError::into_inner)
3938                .active_host_reservations_mb,
3939            1024
3940        );
3941
3942        drop(peer);
3943        drop(detached);
3944        assert_eq!(
3945            coordinator
3946                .state
3947                .lock()
3948                .unwrap_or_else(std::sync::PoisonError::into_inner)
3949                .active_host_reservations_mb,
3950            512
3951        );
3952        drop(request);
3953        assert_eq!(
3954            coordinator
3955                .state
3956                .lock()
3957                .unwrap_or_else(std::sync::PoisonError::into_inner)
3958                .active_host_reservations_mb,
3959            0
3960        );
3961    }
3962
3963    #[test]
3964    fn concurrent_non_aligned_starts_transfer_cold_weights_to_pending_without_double_charge() {
3965        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3966            ResourcePolicy::custom_gb(8.0).unwrap(),
3967            hardware(32 * 1024, GpuBackend::Cpu, None),
3968            Arc::new(FixedLiveMemoryProbe::known(24_000)),
3969        ));
3970        let first_allocation = vllm_process_allocation_id("managed/model-a");
3971        let second_allocation = vllm_process_allocation_id("managed/model-b");
3972        let mib = 1024 * 1024_u64;
3973        let first_measured_bytes = 1024 * mib + 1;
3974        let second_measured_bytes = 512 * mib + 1;
3975        let mut first = coordinator
3976            .reserve_measured_host_allocation(
3977                "managed/model-a",
3978                &first_allocation,
3979                first_measured_bytes,
3980                128,
3981            )
3982            .unwrap();
3983        let mut second = coordinator
3984            .reserve_measured_host_allocation(
3985                "managed/model-b",
3986                &second_allocation,
3987                second_measured_bytes,
3988                256,
3989            )
3990            .unwrap();
3991
3992        first.transfer_cold_weights_to_pending_allocation(&first_allocation, first_measured_bytes);
3993        second
3994            .transfer_cold_weights_to_pending_allocation(&second_allocation, second_measured_bytes);
3995
3996        assert_eq!(
3997            coordinator.resident_model_mb(),
3998            1538,
3999            "each non-MiB-aligned allocation must be rounded up independently"
4000        );
4001        assert_eq!(
4002            coordinator
4003                .state
4004                .lock()
4005                .unwrap_or_else(std::sync::PoisonError::into_inner)
4006                .active_host_reservations_mb,
4007            384,
4008            "pending weights replace cold request weights while request overhead remains active"
4009        );
4010
4011        first.publish_resident_weights_as(&first_allocation, first_measured_bytes);
4012        second.publish_resident_weights_as(&second_allocation, second_measured_bytes);
4013        assert_eq!(coordinator.resident_model_mb(), 1538);
4014        assert_eq!(
4015            coordinator
4016                .state
4017                .lock()
4018                .unwrap_or_else(std::sync::PoisonError::into_inner)
4019                .active_host_reservations_mb,
4020            384,
4021            "publication must not subtract transferred weights twice"
4022        );
4023
4024        drop(first);
4025        drop(second);
4026        assert_eq!(
4027            coordinator
4028                .state
4029                .lock()
4030                .unwrap_or_else(std::sync::PoisonError::into_inner)
4031                .active_host_reservations_mb,
4032            0
4033        );
4034        assert_eq!(coordinator.resident_model_mb(), 1538);
4035    }
4036
4037    #[test]
4038    fn pending_allocation_identity_is_idempotent_and_sibling_exact() {
4039        let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
4040            ResourcePolicy::custom_gb(8.0).unwrap(),
4041            hardware(32 * 1024, GpuBackend::Cpu, None),
4042            Arc::new(FixedLiveMemoryProbe::known(24_000)),
4043        ));
4044        coordinator.mark_teardown_pending_allocation_with_charge(
4045            "same/model",
4046            "worker:same/model",
4047            512 * 1024 * 1024,
4048        );
4049        coordinator.mark_teardown_pending_allocation_with_charge(
4050            "same/model",
4051            "worker:same/model",
4052            512 * 1024 * 1024,
4053        );
4054        coordinator.mark_teardown_pending_allocation_with_charge(
4055            "same/model",
4056            "vllm:same/model",
4057            256 * 1024 * 1024,
4058        );
4059        assert_eq!(coordinator.resident_model_mb(), 768);
4060        coordinator.finish_teardown_allocation("same/model", "worker:same/model");
4061        assert!(coordinator.teardown_pending("same/model"));
4062        assert_eq!(coordinator.resident_model_mb(), 256);
4063        coordinator.finish_teardown_allocation("same/model", "worker:same/model");
4064        assert_eq!(coordinator.resident_model_mb(), 256);
4065        coordinator.finish_teardown_allocation("same/model", "vllm:same/model");
4066        assert!(!coordinator.teardown_pending("same/model"));
4067    }
4068
4069    #[test]
4070    fn scoped_coordinator_identity_is_stable_for_one_state_root() {
4071        let root = tempfile::tempdir().unwrap();
4072        let first = scoped_local_admission(
4073            root.path(),
4074            ResourcePolicy::custom_gb(4.0).unwrap(),
4075            hardware(32 * 1024, GpuBackend::Cpu, None),
4076        );
4077        let second = scoped_local_admission(
4078            root.path(),
4079            ResourcePolicy::custom_gb(8.0).unwrap(),
4080            hardware(32 * 1024, GpuBackend::Cpu, None),
4081        );
4082        assert!(Arc::ptr_eq(&first, &second));
4083        assert_eq!(second.policy(), ResourcePolicy::custom_gb(8.0).unwrap());
4084    }
4085
4086    #[cfg(unix)]
4087    #[test]
4088    fn scoped_coordinator_identity_unifies_symlinked_state_roots() {
4089        use std::os::unix::fs::symlink;
4090
4091        let fixture = tempfile::tempdir().unwrap();
4092        let real = fixture.path().join("real-state");
4093        std::fs::create_dir(&real).unwrap();
4094        let alias = fixture.path().join("state-alias");
4095        symlink(&real, &alias).unwrap();
4096
4097        let first = scoped_local_admission(
4098            &real,
4099            ResourcePolicy::custom_gb(4.0).unwrap(),
4100            hardware(32 * 1024, GpuBackend::Cpu, None),
4101        );
4102        let second = scoped_local_admission(
4103            &alias,
4104            ResourcePolicy::custom_gb(8.0).unwrap(),
4105            hardware(32 * 1024, GpuBackend::Cpu, None),
4106        );
4107        assert!(Arc::ptr_eq(&first, &second));
4108        assert_eq!(second.policy(), ResourcePolicy::custom_gb(8.0).unwrap());
4109    }
4110
4111    #[test]
4112    fn unavailable_state_root_normalization_is_absolute_and_lexically_stable() {
4113        let fixture = tempfile::tempdir().unwrap();
4114        let missing = fixture.path().join("not-created").join("..").join("state");
4115        assert_eq!(
4116            normalized_state_root_key(&missing),
4117            normalized_state_root_key(&fixture.path().join("state"))
4118        );
4119    }
4120
4121    #[cfg(unix)]
4122    #[test]
4123    fn missing_leaf_under_symlinked_parent_keeps_one_scope_identity() {
4124        use std::os::unix::fs::symlink;
4125
4126        let fixture = tempfile::tempdir().unwrap();
4127        let real = fixture.path().join("real");
4128        std::fs::create_dir(&real).unwrap();
4129        let alias = fixture.path().join("alias");
4130        symlink(&real, &alias).unwrap();
4131
4132        assert_eq!(
4133            normalized_state_root_key(&alias.join("missing").join("state")),
4134            normalized_state_root_key(&real.join("missing").join("state"))
4135        );
4136    }
4137
4138    #[cfg(unix)]
4139    #[test]
4140    fn parent_components_are_resolved_after_symlinks_not_lexically_before_them() {
4141        use std::os::unix::fs::symlink;
4142
4143        let fixture = tempfile::tempdir().unwrap();
4144        let physical_parent = fixture.path().join("physical");
4145        let physical_child = physical_parent.join("child");
4146        std::fs::create_dir_all(&physical_child).unwrap();
4147        let aliases = fixture.path().join("aliases");
4148        std::fs::create_dir(&aliases).unwrap();
4149        let alias = aliases.join("runtime");
4150        symlink(&physical_child, &alias).unwrap();
4151
4152        let through_alias = alias.join("..").join("missing-state");
4153        assert_eq!(
4154            normalized_state_root_key(&through_alias),
4155            normalized_state_root_key(&physical_parent.join("missing-state"))
4156        );
4157        assert_ne!(
4158            normalized_state_root_key(&through_alias),
4159            aliases.join("missing-state")
4160        );
4161    }
4162
4163    fn assert_no_temp_files(directory: &std::path::Path) {
4164        let entries = std::fs::read_dir(directory)
4165            .unwrap()
4166            .map(|entry| entry.unwrap().path())
4167            .collect::<Vec<_>>();
4168        assert_eq!(entries, vec![directory.join("model-resource-policy.json")]);
4169    }
4170
4171    #[cfg(unix)]
4172    fn assert_private_mode(path: &std::path::Path, expected: u32) {
4173        use std::os::unix::fs::PermissionsExt;
4174        assert_eq!(
4175            std::fs::metadata(path).unwrap().permissions().mode() & 0o777,
4176            expected
4177        );
4178    }
4179
4180    #[cfg(not(unix))]
4181    fn assert_private_mode(_path: &std::path::Path, _expected: u32) {}
4182}