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