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