1use 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
22pub 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
31pub 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;
39const BYTES_PER_MB: u128 = 1024 * 1024;
41
42#[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#[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 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 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 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 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#[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#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
201pub struct AcceleratorResourceBudget {
202 pub total_mb: u64,
203 pub budget_mb: u64,
204}
205
206#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
208pub struct ResourceEvaluation {
209 pub host_memory: EffectiveResourceBudget,
210 pub accelerator_memory: Option<AcceleratorResourceBudget>,
211}
212
213#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
215#[serde(rename_all = "snake_case")]
216pub enum ModelResourceEvidence {
217 CatalogExact,
219 FileSystemMeasured,
221 Heuristic,
223}
224
225#[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#[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 PendingTeardown,
269}
270
271impl LocalLoadVerdict {
272 pub fn permits_static_fallback(&self) -> bool {
273 matches!(self, Self::Allowed | Self::LiveMemoryUnknown)
274 }
275}
276
277pub 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
358pub 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
369pub 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 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
411pub 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 if let Ok(canonical) = std::fs::canonicalize(&absolute) {
428 return canonical;
429 }
430
431 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 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
489pub 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 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 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 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 pub fn mark_teardown_pending(&self, model_id: &str) {
823 self.mark_teardown_pending_allocation(model_id, model_id);
824 }
825
826 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 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 pub fn finish_teardown(&self, model_id: &str) {
881 self.finish_teardown_allocation(model_id, model_id);
882 }
883
884 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 #[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 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 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 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 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
1500pub 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#[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 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 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 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 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 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 pub fn publish_resident_weights_as(
1897 &mut self,
1898 allocation_id: &str,
1899 measured_weights_bytes: u64,
1900 ) {
1901 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 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 _ => 2,
2038 }
2039}
2040
2041fn 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
2073fn 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
2085fn 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
2095fn 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
2108pub 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
2117pub 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
2222pub 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#[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#[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#[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#[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 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 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 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}