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