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