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