Skip to main content

car_inference/
resource_policy.rs

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