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, (ConfigStamp, 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 geometry = self.kv_geometry(model);
729 estimate_model_memory_with_geometry(
730 model,
731 &self.hardware,
732 context_tokens,
733 None,
734 geometry.as_ref(),
735 )
736 }
737
738 fn kv_geometry(&self, model: &ModelSchema) -> Option<KvGeometry> {
745 let models_dir = self
746 .models_dir
747 .read()
748 .unwrap_or_else(std::sync::PoisonError::into_inner)
749 .clone();
750 let mut cache = self
751 .kv_geometry_cache
752 .lock()
753 .unwrap_or_else(std::sync::PoisonError::into_inner);
754 let Some(stamp) = models_dir
757 .as_deref()
758 .and_then(|dir| ConfigStamp::of(&kv_config_path(model, dir)))
759 else {
760 cache.remove(&model.id);
761 return None;
762 };
763 if let Some((cached, geometry)) = cache.get(&model.id) {
764 if *cached == stamp {
765 return Some(geometry.clone());
766 }
767 }
768 let geometry = kv_geometry_for(model, models_dir.as_deref());
769 match geometry.as_ref() {
770 Some(geometry) => {
771 cache.insert(model.id.clone(), (stamp, geometry.clone()));
772 }
773 None => {
774 cache.remove(&model.id);
775 }
776 }
777 geometry
778 }
779
780 pub fn checkpoint_geometry<'a>(
783 &self,
784 models: impl IntoIterator<Item = &'a ModelSchema>,
785 ) -> CheckpointGeometry {
786 CheckpointGeometry {
787 by_id: models
788 .into_iter()
789 .filter_map(|model| Some((model.id.clone(), self.kv_geometry(model)?)))
790 .collect(),
791 }
792 }
793
794 fn set_models_dir(&self, models_dir: PathBuf) {
795 let mut configured = self
796 .models_dir
797 .write()
798 .unwrap_or_else(std::sync::PoisonError::into_inner);
799 if configured.as_ref() == Some(&models_dir) {
800 return;
801 }
802 *configured = Some(models_dir);
803 self.kv_geometry_cache
804 .lock()
805 .unwrap_or_else(std::sync::PoisonError::into_inner)
806 .clear();
807 }
808
809 pub fn set_policy(&self, policy: ResourcePolicy) {
810 *self
811 .policy
812 .write()
813 .unwrap_or_else(std::sync::PoisonError::into_inner) = policy;
814 }
815
816 pub fn policy(&self) -> ResourcePolicy {
817 self.policy
818 .read()
819 .unwrap_or_else(std::sync::PoisonError::into_inner)
820 .clone()
821 }
822
823 pub fn mark_resident(&self, model_id: &str, weights_mb: u64) {
824 let placement = self.default_weight_placement();
825 self.mark_resident_with_placement(model_id, model_id, weights_mb, placement);
826 }
827
828 pub fn mark_resident_allocation(
829 &self,
830 logical_model_id: &str,
831 allocation_id: &str,
832 weights_mb: u64,
833 ) {
834 let placement = self.default_weight_placement();
835 self.mark_resident_with_placement(logical_model_id, allocation_id, weights_mb, placement);
836 }
837
838 pub fn mark_evicted(&self, model_id: &str) {
839 let mut state = self
840 .state
841 .lock()
842 .unwrap_or_else(std::sync::PoisonError::into_inner);
843 state.resident_models.remove(model_id);
846 self.machine_ledger
847 .lock()
848 .unwrap_or_else(std::sync::PoisonError::into_inner)
849 .resident_models
850 .remove(&(self.owner_id, model_id.to_string()));
851 self.resident_activity_leases
852 .lock()
853 .unwrap_or_else(std::sync::PoisonError::into_inner)
854 .remove(model_id);
855 }
856
857 pub fn mark_teardown_pending(&self, model_id: &str) {
860 self.mark_teardown_pending_allocation(model_id, model_id);
861 }
862
863 pub fn mark_teardown_pending_allocation(&self, logical_model_id: &str, allocation_id: &str) {
867 self.mark_teardown_pending_allocation_with_charge(logical_model_id, allocation_id, 0);
868 }
869
870 pub fn mark_teardown_pending_allocation_with_charge(
875 &self,
876 logical_model_id: &str,
877 allocation_id: &str,
878 measured_bytes: u64,
879 ) {
880 let mut state = self
881 .state
882 .lock()
883 .unwrap_or_else(std::sync::PoisonError::into_inner);
884 let model_id = Self::resolve_model_identity(&state, logical_model_id);
885 state
886 .pending_teardown_models
887 .entry(model_id.clone())
888 .or_default()
889 .insert(allocation_id.to_string());
890 let mut machine = self
891 .machine_ledger
892 .lock()
893 .unwrap_or_else(std::sync::PoisonError::into_inner);
894 let key = (self.owner_id, allocation_id.to_string());
895 let machine_resident = machine.resident_models.remove(&key);
896 let scoped_resident = state.resident_models.get(allocation_id).cloned();
897 let resident = machine_resident.or(scoped_resident);
898 let measured_mb = measured_bytes.div_ceil(1024 * 1024);
899 let placement = self.default_weight_placement();
900 machine
901 .pending_allocations
902 .entry(key)
903 .and_modify(|pending| {
904 pending.weights_mb = pending.weights_mb.max(measured_mb);
905 pending.logical_model_id = model_id.clone();
906 })
907 .or_insert_with(|| {
908 resident.unwrap_or(ResidentAllocation {
909 weights_mb: measured_mb,
910 placement,
911 logical_model_id: model_id,
912 })
913 });
914 }
915
916 pub fn finish_teardown(&self, model_id: &str) {
918 self.finish_teardown_allocation(model_id, model_id);
919 }
920
921 pub fn finish_teardown_allocation(&self, logical_model_id: &str, allocation_id: &str) {
925 let mut state = self
926 .state
927 .lock()
928 .unwrap_or_else(std::sync::PoisonError::into_inner);
929 let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
930 let exact_pending = state
931 .pending_teardown_models
932 .get_mut(&logical_model_id)
933 .is_some_and(|pending| pending.remove(allocation_id));
934 if exact_pending
935 && state
936 .pending_teardown_models
937 .get(&logical_model_id)
938 .is_some_and(HashSet::is_empty)
939 {
940 state.pending_teardown_models.remove(&logical_model_id);
941 }
942 if !exact_pending {
943 return;
944 }
945 state.resident_models.remove(allocation_id);
946 let mut machine = self
947 .machine_ledger
948 .lock()
949 .unwrap_or_else(std::sync::PoisonError::into_inner);
950 machine
951 .resident_models
952 .remove(&(self.owner_id, allocation_id.to_string()));
953 machine
954 .pending_allocations
955 .remove(&(self.owner_id, allocation_id.to_string()));
956 self.resident_activity_leases
957 .lock()
958 .unwrap_or_else(std::sync::PoisonError::into_inner)
959 .remove(allocation_id);
960 }
961
962 pub fn teardown_pending(&self, model_id: &str) -> bool {
963 let state = self
964 .state
965 .lock()
966 .unwrap_or_else(std::sync::PoisonError::into_inner);
967 let model_id = Self::resolve_model_identity(&state, model_id);
968 state.pending_teardown_models.contains_key(&model_id)
969 }
970
971 pub fn is_resident(&self, model_id: &str) -> bool {
972 let state = self
973 .state
974 .lock()
975 .unwrap_or_else(std::sync::PoisonError::into_inner);
976 let model_id = Self::resolve_model_identity(&state, model_id);
977 state
978 .resident_models
979 .iter()
980 .any(|(allocation_id, resident)| {
981 allocation_id == &model_id || resident.logical_model_id == model_id
982 })
983 }
984
985 pub fn resident_model_mb(&self) -> u64 {
986 let machine = self
987 .machine_ledger
988 .lock()
989 .unwrap_or_else(std::sync::PoisonError::into_inner);
990 machine
991 .resident_models
992 .values()
993 .chain(machine.pending_allocations.values())
994 .map(|resident| resident.weights_mb)
995 .fold(0, u64::saturating_add)
996 }
997
998 #[doc(hidden)]
1001 pub fn owned_resident_model_mb_for_testing(&self) -> u64 {
1002 let machine = self
1003 .machine_ledger
1004 .lock()
1005 .unwrap_or_else(std::sync::PoisonError::into_inner);
1006 machine
1007 .resident_models
1008 .iter()
1009 .chain(machine.pending_allocations.iter())
1010 .filter(|((owner_id, _), _)| *owner_id == self.owner_id)
1011 .map(|(_, resident)| resident.weights_mb)
1012 .fold(0, u64::saturating_add)
1013 }
1014
1015 pub fn active_request_count(&self, model_id: &str) -> usize {
1016 let state = self
1017 .state
1018 .lock()
1019 .unwrap_or_else(std::sync::PoisonError::into_inner);
1020 let model_id = Self::resolve_model_identity(&state, model_id);
1021 state
1022 .active_by_model
1023 .get(&model_id)
1024 .copied()
1025 .unwrap_or_default()
1026 }
1027
1028 pub fn resident_allocation_ids(&self, model_id: &str) -> Vec<String> {
1032 let state = self
1033 .state
1034 .lock()
1035 .unwrap_or_else(std::sync::PoisonError::into_inner);
1036 let model_id = Self::resolve_model_identity(&state, model_id);
1037 let mut allocations = state
1038 .resident_models
1039 .iter()
1040 .filter(|(allocation_id, resident)| {
1041 allocation_id.as_str() == model_id || resident.logical_model_id == model_id
1042 })
1043 .map(|(allocation_id, _)| allocation_id.clone())
1044 .collect::<Vec<_>>();
1045 if let Some(pending) = state.pending_teardown_models.get(&model_id) {
1046 allocations.extend(pending.iter().cloned());
1047 }
1048 allocations.sort();
1049 allocations.dedup();
1050 allocations
1051 }
1052
1053 fn default_weight_placement(&self) -> WeightPlacement {
1054 if matches!(self.hardware.gpu_backend, GpuBackend::Cuda)
1055 && self.hardware.gpu_memory_mb.is_some()
1056 {
1057 WeightPlacement::Accelerator
1058 } else {
1059 WeightPlacement::Host
1060 }
1061 }
1062
1063 fn mark_resident_with_placement(
1064 &self,
1065 logical_model_id: &str,
1066 allocation_id: &str,
1067 weights_mb: u64,
1068 placement: WeightPlacement,
1069 ) {
1070 let mut state = self
1071 .state
1072 .lock()
1073 .unwrap_or_else(std::sync::PoisonError::into_inner);
1074 let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
1075 let allocation = ResidentAllocation {
1076 weights_mb,
1077 placement,
1078 logical_model_id: logical_model_id.clone(),
1079 };
1080 state
1081 .resident_models
1082 .entry(allocation_id.to_string())
1083 .and_modify(|resident| {
1084 resident.weights_mb = resident.weights_mb.max(weights_mb);
1085 resident.placement = placement;
1086 resident.logical_model_id = logical_model_id.clone();
1087 })
1088 .or_insert(ResidentAllocation {
1089 weights_mb,
1090 placement,
1091 logical_model_id,
1092 });
1093 self.machine_ledger
1095 .lock()
1096 .unwrap_or_else(std::sync::PoisonError::into_inner)
1097 .resident_models
1098 .insert((self.owner_id, allocation_id.to_string()), allocation);
1099 }
1100
1101 pub fn preflight(&self, model: &ModelSchema, context_tokens: usize) -> LocalLoadPreflight {
1102 let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
1103 let state = self
1104 .state
1105 .lock()
1106 .unwrap_or_else(std::sync::PoisonError::into_inner);
1107 let machine = self
1108 .machine_ledger
1109 .lock()
1110 .unwrap_or_else(std::sync::PoisonError::into_inner);
1111 self.preflight_locked(model, context_tokens, &state, &machine, live_available_mb)
1112 }
1113
1114 pub fn reserve(
1115 self: &Arc<Self>,
1116 model: &ModelSchema,
1117 context_tokens: usize,
1118 ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1119 let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
1120 let mut state = self
1121 .state
1122 .lock()
1123 .unwrap_or_else(std::sync::PoisonError::into_inner);
1124 let mut machine = self
1125 .machine_ledger
1126 .lock()
1127 .unwrap_or_else(std::sync::PoisonError::into_inner);
1128 let model_id = Self::resolve_model_identity(&state, &model.id);
1129 let estimate = self.model_memory_estimate(model, context_tokens);
1130 let mut preflight = self.preflight_estimate_locked(
1131 &model_id,
1132 estimate,
1133 self.default_weight_placement(),
1134 &state,
1135 &machine,
1136 live_available_mb,
1137 );
1138 if state.pending_teardown_models.contains_key(&model_id) {
1139 preflight.verdict = LocalLoadVerdict::PendingTeardown;
1140 return Err(LocalAdmissionError { preflight });
1141 }
1142 if state.maintenance_models.contains(&model_id) {
1143 preflight.verdict = LocalLoadVerdict::ModelMaintenance;
1144 return Err(LocalAdmissionError { preflight });
1145 }
1146 if !preflight.verdict.permits_static_fallback() {
1147 return Err(LocalAdmissionError { preflight });
1148 }
1149 if preflight.verdict == LocalLoadVerdict::LiveMemoryUnknown {
1150 tracing::warn!(
1151 model = %model.id,
1152 configured_ceiling_mb = preflight.configured_ceiling_mb,
1153 estimated_incremental_mb = preflight.estimated_incremental_mb,
1154 "live memory is unknown; proceeding under the static configured ceiling only"
1155 );
1156 }
1157 state.active_host_reservations_mb = state
1158 .active_host_reservations_mb
1159 .saturating_add(preflight.estimated_incremental_mb);
1160 state.active_accelerator_reservations_mb = state
1161 .active_accelerator_reservations_mb
1162 .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1163 let host = machine
1164 .active_host_by_owner
1165 .entry(self.owner_id)
1166 .or_default();
1167 *host = host.saturating_add(preflight.estimated_incremental_mb);
1168 let accelerator = machine
1169 .active_accelerator_by_owner
1170 .entry(self.owner_id)
1171 .or_default();
1172 *accelerator =
1173 accelerator.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1174 *state.active_by_model.entry(model_id.clone()).or_default() += 1;
1175 state.next_request_id = state.next_request_id.wrapping_add(1);
1176 let request_id = state.next_request_id;
1177 Ok(LocalLoadReservation {
1178 request_id,
1179 model_id: model_id.clone(),
1180 maintenance_model_id: model_id.clone(),
1181 weights_mb: preflight.estimate.weights_mb,
1182 reserved_incremental_mb: preflight.estimated_incremental_mb,
1183 reserved_accelerator_mb: preflight.accelerator_incremental_mb.unwrap_or_default(),
1184 cold_weights_reserved: !Self::has_resident_model(&state, &model_id),
1185 placement: self.default_weight_placement(),
1186 coordinator: Arc::clone(self),
1187 charge: Arc::new(ReservationCharge::new(
1188 model_id,
1189 preflight.estimated_incremental_mb,
1190 preflight.accelerator_incremental_mb.unwrap_or_default(),
1191 Arc::clone(self),
1192 )),
1193 })
1194 }
1195
1196 pub fn reserve_measured_host(
1200 self: &Arc<Self>,
1201 model_id: &str,
1202 measured_weights_bytes: u64,
1203 request_overhead_mb: u64,
1204 ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1205 self.reserve_measured(
1206 model_id,
1207 model_id,
1208 measured_weights_bytes,
1209 request_overhead_mb,
1210 WeightPlacement::Host,
1211 )
1212 }
1213
1214 pub fn reserve_measured_host_allocation(
1215 self: &Arc<Self>,
1216 logical_model_id: &str,
1217 allocation_id: &str,
1218 measured_weights_bytes: u64,
1219 request_overhead_mb: u64,
1220 ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1221 self.reserve_measured(
1222 logical_model_id,
1223 allocation_id,
1224 measured_weights_bytes,
1225 request_overhead_mb,
1226 WeightPlacement::Host,
1227 )
1228 }
1229
1230 fn reserve_measured(
1231 self: &Arc<Self>,
1232 logical_model_id: &str,
1233 allocation_id: &str,
1234 measured_weights_bytes: u64,
1235 request_overhead_mb: u64,
1236 placement: WeightPlacement,
1237 ) -> Result<LocalLoadReservation, LocalAdmissionError> {
1238 let weights_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1239 let estimate = ModelMemoryEstimate {
1240 weights_mb,
1241 runtime_overhead_mb: request_overhead_mb,
1242 context_overhead_mb: 0,
1243 transient_margin_mb: 0,
1244 estimated_peak_mb: weights_mb.saturating_add(request_overhead_mb),
1245 evidence: ModelResourceEvidence::FileSystemMeasured,
1246 };
1247 let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
1248 let mut state = self
1249 .state
1250 .lock()
1251 .unwrap_or_else(std::sync::PoisonError::into_inner);
1252 let mut machine = self
1253 .machine_ledger
1254 .lock()
1255 .unwrap_or_else(std::sync::PoisonError::into_inner);
1256 let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
1257 let mut preflight = self.preflight_estimate_locked(
1258 allocation_id,
1259 estimate,
1260 placement,
1261 &state,
1262 &machine,
1263 live_available_mb,
1264 );
1265 if state
1266 .pending_teardown_models
1267 .contains_key(&logical_model_id)
1268 {
1269 preflight.verdict = LocalLoadVerdict::PendingTeardown;
1270 }
1271 if state.maintenance_models.contains(&logical_model_id) {
1272 preflight.verdict = LocalLoadVerdict::ModelMaintenance;
1273 }
1274 if !preflight.verdict.permits_static_fallback() {
1275 return Err(LocalAdmissionError { preflight });
1276 }
1277 state.active_host_reservations_mb = state
1278 .active_host_reservations_mb
1279 .saturating_add(preflight.estimated_incremental_mb);
1280 state.active_accelerator_reservations_mb = state
1281 .active_accelerator_reservations_mb
1282 .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1283 let host = machine
1284 .active_host_by_owner
1285 .entry(self.owner_id)
1286 .or_default();
1287 *host = host.saturating_add(preflight.estimated_incremental_mb);
1288 let accelerator = machine
1289 .active_accelerator_by_owner
1290 .entry(self.owner_id)
1291 .or_default();
1292 *accelerator =
1293 accelerator.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1294 *state
1295 .active_by_model
1296 .entry(logical_model_id.clone())
1297 .or_default() += 1;
1298 state.next_request_id = state.next_request_id.wrapping_add(1);
1299 Ok(LocalLoadReservation {
1300 request_id: state.next_request_id,
1301 model_id: allocation_id.to_string(),
1302 maintenance_model_id: logical_model_id.clone(),
1303 weights_mb,
1304 reserved_incremental_mb: preflight.estimated_incremental_mb,
1305 reserved_accelerator_mb: preflight.accelerator_incremental_mb.unwrap_or_default(),
1306 cold_weights_reserved: !state.resident_models.contains_key(allocation_id),
1307 placement,
1308 coordinator: Arc::clone(self),
1309 charge: Arc::new(ReservationCharge::new(
1310 logical_model_id,
1311 preflight.estimated_incremental_mb,
1312 preflight.accelerator_incremental_mb.unwrap_or_default(),
1313 Arc::clone(self),
1314 )),
1315 })
1316 }
1317
1318 pub fn begin_model_maintenance(
1322 self: &Arc<Self>,
1323 model_id: &str,
1324 ) -> Result<LocalModelMaintenanceGuard, ModelMaintenanceError> {
1325 let mut state = self
1326 .state
1327 .lock()
1328 .unwrap_or_else(std::sync::PoisonError::into_inner);
1329 let model_id = Self::resolve_model_identity(&state, model_id);
1330 if state
1331 .active_by_model
1332 .get(&model_id)
1333 .copied()
1334 .unwrap_or_default()
1335 > 0
1336 {
1337 return Err(ModelMaintenanceError::ModelInUse(model_id));
1338 }
1339 if !state.maintenance_models.insert(model_id.clone()) {
1340 return Err(ModelMaintenanceError::AlreadyInMaintenance(model_id));
1341 }
1342 Ok(LocalModelMaintenanceGuard {
1343 model_id,
1344 coordinator: Arc::clone(self),
1345 })
1346 }
1347
1348 fn preflight_locked(
1349 &self,
1350 model: &ModelSchema,
1351 context_tokens: usize,
1352 state: &AdmissionState,
1353 machine: &MachineAdmissionLedger,
1354 live_available_mb: Option<u64>,
1355 ) -> LocalLoadPreflight {
1356 let estimate = self.model_memory_estimate(model, context_tokens);
1357 let model_id = Self::resolve_model_identity(state, &model.id);
1358 self.preflight_estimate_locked(
1359 &model_id,
1360 estimate,
1361 self.default_weight_placement(),
1362 state,
1363 machine,
1364 live_available_mb,
1365 )
1366 }
1367
1368 fn preflight_estimate_locked(
1369 &self,
1370 model_id: &str,
1371 estimate: ModelMemoryEstimate,
1372 placement: WeightPlacement,
1373 state: &AdmissionState,
1374 machine: &MachineAdmissionLedger,
1375 live_available_mb: Option<u64>,
1376 ) -> LocalLoadPreflight {
1377 let policy = self.policy();
1378 let budget = policy.effective_budget(self.hardware.total_ram_mb);
1379 let host_resident_mb = machine
1380 .resident_models
1381 .values()
1382 .chain(machine.pending_allocations.values())
1383 .filter(|resident| resident.placement == WeightPlacement::Host)
1384 .map(|resident| resident.weights_mb)
1385 .fold(0, u64::saturating_add);
1386 let accelerator_resident_weights_mb = machine
1387 .resident_models
1388 .values()
1389 .chain(machine.pending_allocations.values())
1390 .filter(|resident| resident.placement == WeightPlacement::Accelerator)
1391 .map(|resident| resident.weights_mb)
1392 .fold(0, u64::saturating_add);
1393 let already_resident = Self::has_resident_model(state, model_id);
1394 let request_overhead_mb = estimate
1395 .context_overhead_mb
1396 .saturating_add(estimate.runtime_overhead_mb)
1397 .saturating_add(estimate.transient_margin_mb);
1398 let cold_weights_mb = if already_resident {
1399 0
1400 } else {
1401 estimate.weights_mb
1402 };
1403 let resident_model_mb = host_resident_mb;
1404 let host_cold_weights_mb = if placement == WeightPlacement::Host {
1405 cold_weights_mb
1406 } else {
1407 0
1408 };
1409 let estimated_incremental_mb = request_overhead_mb.saturating_add(host_cold_weights_mb);
1410 let accelerator_total_mb = (placement == WeightPlacement::Accelerator)
1411 .then_some(self.hardware.gpu_memory_mb)
1412 .flatten();
1413 let accelerator_resident_mb =
1414 (placement == WeightPlacement::Accelerator).then_some(accelerator_resident_weights_mb);
1415 let accelerator_incremental_mb =
1416 (placement == WeightPlacement::Accelerator).then_some(cold_weights_mb);
1417 let projected_static_mb = resident_model_mb
1418 .saturating_add(machine.active_host_by_owner.values().copied().sum::<u64>())
1419 .saturating_add(estimated_incremental_mb);
1420 let projected_accelerator_mb = accelerator_resident_weights_mb
1421 .saturating_add(
1422 machine
1423 .active_accelerator_by_owner
1424 .values()
1425 .copied()
1426 .sum::<u64>(),
1427 )
1428 .saturating_add(cold_weights_mb);
1429
1430 let verdict = if budget.effective_new_load_ceiling_mb == 0 && !already_resident {
1431 LocalLoadVerdict::DisabledByPolicy
1432 } else if (budget.configured_model_ceiling_mb != 0
1433 && projected_static_mb > budget.configured_model_ceiling_mb)
1434 || accelerator_total_mb.is_some_and(|vram_mb| projected_accelerator_mb > vram_mb)
1435 {
1436 LocalLoadVerdict::ExceedsConfiguredCeiling
1437 } else if let Some(available_mb) = live_available_mb {
1438 let unreserved_available_mb = available_mb
1439 .saturating_sub(machine.active_host_by_owner.values().copied().sum::<u64>());
1440 if unreserved_available_mb
1441 < estimated_incremental_mb.saturating_add(budget.emergency_reserve_mb)
1442 {
1443 LocalLoadVerdict::InsufficientLiveMemory
1444 } else {
1445 LocalLoadVerdict::Allowed
1446 }
1447 } else {
1448 LocalLoadVerdict::LiveMemoryUnknown
1449 };
1450
1451 LocalLoadPreflight {
1452 model_id: model_id.to_string(),
1453 estimate,
1454 configured_ceiling_mb: budget.configured_model_ceiling_mb,
1455 resident_model_mb,
1456 active_reservations_mb: machine.active_host_by_owner.values().copied().sum(),
1457 estimated_incremental_mb,
1458 accelerator_total_mb,
1459 accelerator_resident_mb,
1460 accelerator_incremental_mb,
1461 live_available_mb,
1462 emergency_reserve_mb: budget.emergency_reserve_mb,
1463 verdict,
1464 }
1465 }
1466
1467 fn has_resident_model(state: &AdmissionState, model_id: &str) -> bool {
1468 state
1469 .resident_models
1470 .iter()
1471 .any(|(allocation_id, resident)| {
1472 allocation_id == model_id || resident.logical_model_id == model_id
1473 })
1474 }
1475}
1476
1477#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
1478pub struct LocalAdmissionError {
1479 pub preflight: LocalLoadPreflight,
1480}
1481
1482#[derive(Clone, Debug, Error, PartialEq, Eq)]
1483pub enum ModelMaintenanceError {
1484 #[error("local model '{0}' is in use")]
1485 ModelInUse(String),
1486 #[error("local model '{0}' already has maintenance in progress")]
1487 AlreadyInMaintenance(String),
1488 #[error("failed to release local model residency: {0}")]
1489 ReleaseFailed(String),
1490 #[error("local worker did not acknowledge release of model '{0}'")]
1491 WorkerReleaseUnacknowledged(String),
1492 #[error("supervised local process did not acknowledge release of model '{0}'")]
1493 ProcessReleaseUnacknowledged(String),
1494 #[error("an in-process cache still has active work for local model '{0}'")]
1495 CacheReleaseBlocked(String),
1496 #[error("local model '{model_id}' still has resident allocations: {allocation_ids:?}")]
1497 ResidualResidency {
1498 model_id: String,
1499 allocation_ids: Vec<String>,
1500 },
1501}
1502
1503pub struct LocalModelMaintenanceGuard {
1504 model_id: String,
1505 coordinator: Arc<LocalAdmissionCoordinator>,
1506}
1507
1508impl Drop for LocalModelMaintenanceGuard {
1509 fn drop(&mut self) {
1510 self.coordinator
1511 .state
1512 .lock()
1513 .unwrap_or_else(std::sync::PoisonError::into_inner)
1514 .maintenance_models
1515 .remove(&self.model_id);
1516 }
1517}
1518
1519impl std::fmt::Display for LocalAdmissionError {
1520 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1521 write!(
1522 formatter,
1523 "local model '{}' blocked by {:?}: needs {} MB incremental, {} MB live available",
1524 self.preflight.model_id,
1525 self.preflight.verdict,
1526 self.preflight.estimated_incremental_mb,
1527 self.preflight
1528 .live_available_mb
1529 .map(|value| value.to_string())
1530 .unwrap_or_else(|| "unknown".into())
1531 )
1532 }
1533}
1534
1535impl std::error::Error for LocalAdmissionError {}
1536
1537pub struct LocalLoadReservation {
1540 request_id: u64,
1541 model_id: String,
1542 maintenance_model_id: String,
1543 weights_mb: u64,
1544 reserved_incremental_mb: u64,
1545 reserved_accelerator_mb: u64,
1546 cold_weights_reserved: bool,
1547 placement: WeightPlacement,
1548 coordinator: Arc<LocalAdmissionCoordinator>,
1549 charge: Arc<ReservationCharge>,
1550}
1551
1552#[derive(Clone)]
1557pub struct DetachedLocalLease {
1558 _charge: Arc<ReservationCharge>,
1559}
1560
1561struct ReservationCharge {
1562 maintenance_model_id: String,
1563 host_mb: AtomicU64,
1564 accelerator_mb: AtomicU64,
1565 cold_weights_transferred: AtomicBool,
1566 coordinator: Arc<LocalAdmissionCoordinator>,
1567 activity_lease: Mutex<Option<Arc<crate::model_management::ModelLease>>>,
1568}
1569
1570impl ReservationCharge {
1571 fn new(
1572 maintenance_model_id: String,
1573 host_mb: u64,
1574 accelerator_mb: u64,
1575 coordinator: Arc<LocalAdmissionCoordinator>,
1576 ) -> Self {
1577 Self {
1578 maintenance_model_id,
1579 host_mb: AtomicU64::new(host_mb),
1580 accelerator_mb: AtomicU64::new(accelerator_mb),
1581 cold_weights_transferred: AtomicBool::new(false),
1582 coordinator,
1583 activity_lease: Mutex::new(None),
1584 }
1585 }
1586
1587 fn update(&self, host_mb: u64, accelerator_mb: u64) {
1588 self.host_mb.store(host_mb, Ordering::Release);
1589 self.accelerator_mb.store(accelerator_mb, Ordering::Release);
1590 }
1591}
1592
1593impl Drop for ReservationCharge {
1594 fn drop(&mut self) {
1595 let host_mb = self.host_mb.load(Ordering::Acquire);
1596 let accelerator_mb = self.accelerator_mb.load(Ordering::Acquire);
1597 let mut state = self
1598 .coordinator
1599 .state
1600 .lock()
1601 .unwrap_or_else(std::sync::PoisonError::into_inner);
1602 let mut machine = self
1603 .coordinator
1604 .machine_ledger
1605 .lock()
1606 .unwrap_or_else(std::sync::PoisonError::into_inner);
1607 state.active_host_reservations_mb =
1608 state.active_host_reservations_mb.saturating_sub(host_mb);
1609 state.active_accelerator_reservations_mb = state
1610 .active_accelerator_reservations_mb
1611 .saturating_sub(accelerator_mb);
1612 if let Some(host) = machine
1613 .active_host_by_owner
1614 .get_mut(&self.coordinator.owner_id)
1615 {
1616 *host = host.saturating_sub(host_mb);
1617 }
1618 if let Some(accelerator) = machine
1619 .active_accelerator_by_owner
1620 .get_mut(&self.coordinator.owner_id)
1621 {
1622 *accelerator = accelerator.saturating_sub(accelerator_mb);
1623 }
1624 let maintenance_model_id =
1625 LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
1626 if let Some(active) = state.active_by_model.get_mut(&maintenance_model_id) {
1627 *active = active.saturating_sub(1);
1628 if *active == 0 {
1629 state.active_by_model.remove(&maintenance_model_id);
1630 }
1631 }
1632 }
1633}
1634
1635impl std::fmt::Debug for LocalLoadReservation {
1636 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1637 formatter
1638 .debug_struct("LocalLoadReservation")
1639 .field("request_id", &self.request_id)
1640 .field("model_id", &self.model_id)
1641 .field("reserved_incremental_mb", &self.reserved_incremental_mb)
1642 .field("reserved_accelerator_mb", &self.reserved_accelerator_mb)
1643 .finish_non_exhaustive()
1644 }
1645}
1646
1647impl LocalLoadReservation {
1648 pub fn request_id(&self) -> u64 {
1649 self.request_id
1650 }
1651
1652 pub fn model_id(&self) -> &str {
1656 &self.model_id
1657 }
1658
1659 pub fn authorizes_model(&self, model_id: &str) -> bool {
1660 self.model_id == model_id || self.maintenance_model_id == model_id
1661 }
1662
1663 pub(crate) fn bind_allocation_id(&mut self, allocation_id: &str) {
1667 self.model_id = allocation_id.to_string();
1668 }
1669
1670 pub fn reserved_incremental_mb(&self) -> u64 {
1671 self.reserved_incremental_mb
1672 }
1673
1674 pub fn reconciled_weights_bytes(&self) -> u64 {
1675 self.weights_mb.saturating_mul(1024 * 1024)
1676 }
1677
1678 pub fn detached_lease(&self) -> DetachedLocalLease {
1679 DetachedLocalLease {
1680 _charge: self.charge.clone(),
1681 }
1682 }
1683
1684 pub(crate) fn attach_activity_lease(&mut self, lease: crate::model_management::ModelLease) {
1685 *self
1686 .charge
1687 .activity_lease
1688 .lock()
1689 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::new(lease));
1690 }
1691
1692 pub(crate) fn transfer_cold_weights_to_pending_allocation(
1696 &self,
1697 allocation_id: &str,
1698 measured_weights_bytes: u64,
1699 ) {
1700 let mut state = self
1701 .coordinator
1702 .state
1703 .lock()
1704 .unwrap_or_else(std::sync::PoisonError::into_inner);
1705 let logical_model_id =
1706 LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
1707 state
1708 .pending_teardown_models
1709 .entry(logical_model_id.clone())
1710 .or_default()
1711 .insert(allocation_id.to_string());
1712 let mut machine = self
1713 .coordinator
1714 .machine_ledger
1715 .lock()
1716 .unwrap_or_else(std::sync::PoisonError::into_inner);
1717 let key = (self.coordinator.owner_id, allocation_id.to_string());
1718 let machine_resident = machine.resident_models.remove(&key);
1719 let scoped_resident = state.resident_models.get(allocation_id).cloned();
1720 let resident = machine_resident.or(scoped_resident);
1721 let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1722 machine
1723 .pending_allocations
1724 .entry(key)
1725 .and_modify(|pending| {
1726 pending.weights_mb = pending.weights_mb.max(measured_mb);
1727 pending.logical_model_id = logical_model_id.clone();
1728 })
1729 .or_insert_with(|| {
1730 resident.unwrap_or(ResidentAllocation {
1731 weights_mb: measured_mb,
1732 placement: self.placement,
1733 logical_model_id,
1734 })
1735 });
1736
1737 if self.cold_weights_reserved
1738 && self
1739 .charge
1740 .cold_weights_transferred
1741 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
1742 .is_ok()
1743 {
1744 match self.placement {
1745 WeightPlacement::Host => {
1746 state.active_host_reservations_mb = state
1747 .active_host_reservations_mb
1748 .saturating_sub(self.weights_mb);
1749 let active = machine
1750 .active_host_by_owner
1751 .entry(self.coordinator.owner_id)
1752 .or_default();
1753 *active = active.saturating_sub(self.weights_mb);
1754 let _ = self.charge.host_mb.fetch_update(
1755 Ordering::AcqRel,
1756 Ordering::Acquire,
1757 |host_mb| Some(host_mb.saturating_sub(self.weights_mb)),
1758 );
1759 }
1760 WeightPlacement::Accelerator => {
1761 state.active_accelerator_reservations_mb = state
1762 .active_accelerator_reservations_mb
1763 .saturating_sub(self.weights_mb);
1764 let active = machine
1765 .active_accelerator_by_owner
1766 .entry(self.coordinator.owner_id)
1767 .or_default();
1768 *active = active.saturating_sub(self.weights_mb);
1769 let _ = self.charge.accelerator_mb.fetch_update(
1770 Ordering::AcqRel,
1771 Ordering::Acquire,
1772 |accelerator_mb| Some(accelerator_mb.saturating_sub(self.weights_mb)),
1773 );
1774 }
1775 }
1776 }
1777 }
1778
1779 fn sync_shared_charge(&self) {
1780 self.charge
1781 .update(self.reserved_incremental_mb, self.reserved_accelerator_mb);
1782 }
1783
1784 pub fn reconcile_measured_weights(
1788 &mut self,
1789 measured_weights_bytes: u64,
1790 ) -> Result<LocalLoadPreflight, LocalAdmissionError> {
1791 if !self.cold_weights_reserved {
1792 let live_available_mb = self
1793 .coordinator
1794 .live_probe
1795 .available_memory_mb()
1796 .ok()
1797 .flatten();
1798 let state = self
1799 .coordinator
1800 .state
1801 .lock()
1802 .unwrap_or_else(std::sync::PoisonError::into_inner);
1803 let machine = self
1804 .coordinator
1805 .machine_ledger
1806 .lock()
1807 .unwrap_or_else(std::sync::PoisonError::into_inner);
1808 if LocalAdmissionCoordinator::has_resident_model(&state, &self.model_id) {
1809 let estimate = self.measured_estimate(self.weights_mb);
1810 return Ok(self.coordinator.preflight_estimate_locked(
1811 &self.model_id,
1812 estimate,
1813 self.placement,
1814 &state,
1815 &machine,
1816 live_available_mb,
1817 ));
1818 }
1819 drop(state);
1820 self.cold_weights_reserved = true;
1821 }
1822 let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1823 let live_available_mb = self
1824 .coordinator
1825 .live_probe
1826 .available_memory_mb()
1827 .ok()
1828 .flatten();
1829 let mut state = self
1830 .coordinator
1831 .state
1832 .lock()
1833 .unwrap_or_else(std::sync::PoisonError::into_inner);
1834 let mut machine = self
1835 .coordinator
1836 .machine_ledger
1837 .lock()
1838 .unwrap_or_else(std::sync::PoisonError::into_inner);
1839 let published_by_peer =
1840 LocalAdmissionCoordinator::has_resident_model(&state, &self.model_id);
1841 let mut without_current = state.clone();
1842 without_current.active_host_reservations_mb = without_current
1843 .active_host_reservations_mb
1844 .saturating_sub(self.reserved_incremental_mb);
1845 without_current.active_accelerator_reservations_mb = without_current
1846 .active_accelerator_reservations_mb
1847 .saturating_sub(self.reserved_accelerator_mb);
1848 let mut without_machine = machine.clone();
1849 let host = without_machine
1850 .active_host_by_owner
1851 .entry(self.coordinator.owner_id)
1852 .or_default();
1853 *host = host.saturating_sub(self.reserved_incremental_mb);
1854 let accelerator = without_machine
1855 .active_accelerator_by_owner
1856 .entry(self.coordinator.owner_id)
1857 .or_default();
1858 *accelerator = accelerator.saturating_sub(self.reserved_accelerator_mb);
1859 let estimate = self.measured_estimate(measured_mb);
1860 let preflight = self.coordinator.preflight_estimate_locked(
1861 &self.model_id,
1862 estimate,
1863 self.placement,
1864 &without_current,
1865 &without_machine,
1866 live_available_mb,
1867 );
1868 if !preflight.verdict.permits_static_fallback() {
1869 return Err(LocalAdmissionError { preflight });
1870 }
1871 state.active_host_reservations_mb = without_current
1872 .active_host_reservations_mb
1873 .saturating_add(preflight.estimated_incremental_mb);
1874 state.active_accelerator_reservations_mb = without_current
1875 .active_accelerator_reservations_mb
1876 .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
1877 machine.active_host_by_owner.insert(
1878 self.coordinator.owner_id,
1879 without_machine
1880 .active_host_by_owner
1881 .get(&self.coordinator.owner_id)
1882 .copied()
1883 .unwrap_or_default()
1884 .saturating_add(preflight.estimated_incremental_mb),
1885 );
1886 machine.active_accelerator_by_owner.insert(
1887 self.coordinator.owner_id,
1888 without_machine
1889 .active_accelerator_by_owner
1890 .get(&self.coordinator.owner_id)
1891 .copied()
1892 .unwrap_or_default()
1893 .saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default()),
1894 );
1895 self.weights_mb = measured_mb;
1896 self.reserved_incremental_mb = preflight.estimated_incremental_mb;
1897 self.reserved_accelerator_mb = preflight.accelerator_incremental_mb.unwrap_or_default();
1898 if published_by_peer {
1899 self.cold_weights_reserved = false;
1900 }
1901 self.sync_shared_charge();
1902 Ok(preflight)
1903 }
1904
1905 fn measured_estimate(&self, weights_mb: u64) -> ModelMemoryEstimate {
1906 let host_cold_mb = if self.cold_weights_reserved && self.placement == WeightPlacement::Host
1907 {
1908 self.weights_mb
1909 } else {
1910 0
1911 };
1912 let request_overhead_mb = self.reserved_incremental_mb.saturating_sub(host_cold_mb);
1913 ModelMemoryEstimate {
1914 weights_mb,
1915 runtime_overhead_mb: request_overhead_mb,
1916 context_overhead_mb: 0,
1917 transient_margin_mb: 0,
1918 estimated_peak_mb: weights_mb.saturating_add(request_overhead_mb),
1919 evidence: ModelResourceEvidence::FileSystemMeasured,
1920 }
1921 }
1922
1923 pub fn publish_resident_weights(&mut self, measured_weights_bytes: u64) {
1926 let allocation_id = self.model_id.clone();
1927 self.publish_resident_weights_as(&allocation_id, measured_weights_bytes);
1928 }
1929
1930 pub fn publish_resident_weights_as(
1934 &mut self,
1935 allocation_id: &str,
1936 measured_weights_bytes: u64,
1937 ) {
1938 let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
1945 let mut state = self
1946 .coordinator
1947 .state
1948 .lock()
1949 .unwrap_or_else(std::sync::PoisonError::into_inner);
1950 let mut machine = self
1951 .coordinator
1952 .machine_ledger
1953 .lock()
1954 .unwrap_or_else(std::sync::PoisonError::into_inner);
1955 let key = (self.coordinator.owner_id, allocation_id.to_string());
1956 let pending = machine.pending_allocations.remove(&key);
1957 let logical_model_id =
1958 LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
1959 if let Some(allocations) = state.pending_teardown_models.get_mut(&logical_model_id) {
1960 allocations.remove(allocation_id);
1961 if allocations.is_empty() {
1962 state.pending_teardown_models.remove(&logical_model_id);
1963 }
1964 }
1965 if !self.cold_weights_reserved && pending.is_none() {
1966 return;
1967 }
1968 if self.cold_weights_reserved {
1969 let cold_weights_transferred =
1970 self.charge.cold_weights_transferred.load(Ordering::Acquire);
1971 match self.placement {
1972 WeightPlacement::Host => {
1973 if !cold_weights_transferred {
1974 state.active_host_reservations_mb = state
1975 .active_host_reservations_mb
1976 .saturating_sub(self.weights_mb);
1977 let active = machine
1978 .active_host_by_owner
1979 .entry(self.coordinator.owner_id)
1980 .or_default();
1981 *active = active.saturating_sub(self.weights_mb);
1982 }
1983 self.reserved_incremental_mb =
1984 self.reserved_incremental_mb.saturating_sub(self.weights_mb);
1985 }
1986 WeightPlacement::Accelerator => {
1987 if !cold_weights_transferred {
1988 state.active_accelerator_reservations_mb = state
1989 .active_accelerator_reservations_mb
1990 .saturating_sub(self.weights_mb);
1991 let active = machine
1992 .active_accelerator_by_owner
1993 .entry(self.coordinator.owner_id)
1994 .or_default();
1995 *active = active.saturating_sub(self.weights_mb);
1996 }
1997 self.reserved_accelerator_mb =
1998 self.reserved_accelerator_mb.saturating_sub(self.weights_mb);
1999 }
2000 }
2001 }
2002 let resident = ResidentAllocation {
2003 weights_mb: pending
2004 .as_ref()
2005 .map(|allocation| allocation.weights_mb)
2006 .unwrap_or_default()
2007 .max(measured_mb),
2008 placement: pending
2009 .as_ref()
2010 .map(|allocation| allocation.placement)
2011 .unwrap_or(self.placement),
2012 logical_model_id,
2013 };
2014 state
2015 .resident_models
2016 .insert(allocation_id.to_string(), resident.clone());
2017 machine.resident_models.insert(key, resident.clone());
2018 if let Some(lease) = self
2019 .charge
2020 .activity_lease
2021 .lock()
2022 .unwrap_or_else(std::sync::PoisonError::into_inner)
2023 .clone()
2024 {
2025 self.coordinator
2026 .resident_activity_leases
2027 .lock()
2028 .unwrap_or_else(std::sync::PoisonError::into_inner)
2029 .insert(allocation_id.to_string(), lease);
2030 }
2031 self.weights_mb = resident.weights_mb;
2032 self.cold_weights_reserved = false;
2033 self.charge
2034 .cold_weights_transferred
2035 .store(false, Ordering::Release);
2036 self.sync_shared_charge();
2037 }
2038
2039 pub fn commit_resident_weights(&mut self) {
2042 self.publish_resident_weights(self.weights_mb * 1024 * 1024);
2043 }
2044}
2045
2046#[derive(Clone, Debug, PartialEq, Eq)]
2047struct KvGeometry {
2048 num_hidden_layers: usize,
2049 num_key_value_heads: usize,
2050 head_dim: usize,
2051 kv_bytes_per_element: usize,
2052 hidden_size: Option<usize>,
2053}
2054
2055#[derive(Clone, Debug, Default)]
2062pub struct CheckpointGeometry {
2063 by_id: HashMap<String, KvGeometry>,
2064}
2065
2066impl CheckpointGeometry {
2067 pub fn read<'a>(models: impl IntoIterator<Item = &'a ModelSchema>, models_dir: &Path) -> Self {
2072 Self {
2073 by_id: models
2074 .into_iter()
2075 .filter_map(|model| {
2076 Some((model.id.clone(), kv_geometry_for(model, Some(models_dir))?))
2077 })
2078 .collect(),
2079 }
2080 }
2081
2082 fn get(&self, model: &ModelSchema) -> Option<&KvGeometry> {
2083 self.by_id.get(&model.id)
2084 }
2085}
2086
2087pub(crate) const NESTED_TEXT_CONFIGS: [&str; 3] = ["text_config", "llm_config", "language_config"];
2092
2093#[derive(Deserialize)]
2094struct ModelGeometryConfig {
2095 num_hidden_layers: Option<usize>,
2096 num_key_value_heads: Option<usize>,
2097 head_dim: Option<usize>,
2098 hidden_size: Option<usize>,
2099 num_attention_heads: Option<usize>,
2100 kv_cache_dtype: Option<String>,
2101 torch_dtype: Option<String>,
2102}
2103
2104fn dtype_bytes(dtype: Option<&str>) -> usize {
2105 match dtype.map(str::to_ascii_lowercase).as_deref() {
2106 Some("float32" | "fp32" | "f32") => 4,
2107 Some("float64" | "fp64" | "f64") => 8,
2108 Some("float8" | "fp8" | "int8" | "uint8" | "i8" | "u8") => 1,
2109 _ => 2,
2113 }
2114}
2115
2116#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2119struct ConfigStamp {
2120 modified: Option<std::time::SystemTime>,
2121 len: u64,
2122}
2123
2124impl ConfigStamp {
2125 fn of(path: &Path) -> Option<Self> {
2126 let metadata = std::fs::metadata(path).ok()?;
2127 Some(Self {
2128 modified: metadata.modified().ok(),
2129 len: metadata.len(),
2130 })
2131 }
2132}
2133
2134fn kv_config_path(model: &ModelSchema, models_dir: &Path) -> PathBuf {
2135 models_dir.join(&model.name).join("config.json")
2136}
2137
2138fn kv_geometry_for(model: &ModelSchema, models_dir: Option<&Path>) -> Option<KvGeometry> {
2149 let bytes = std::fs::read(kv_config_path(model, models_dir?)).ok()?;
2150 let raw: serde_json::Value = serde_json::from_slice(&bytes).ok()?;
2151 let source = if raw.get("num_hidden_layers").is_some() {
2152 &raw
2153 } else {
2154 NESTED_TEXT_CONFIGS
2155 .iter()
2156 .find_map(|key| {
2157 raw.get(*key)
2158 .filter(|v| v.get("num_hidden_layers").is_some())
2159 })
2160 .unwrap_or(&raw)
2161 };
2162 let mut config: ModelGeometryConfig = serde_json::from_value(source.clone()).ok()?;
2163 if config.kv_cache_dtype.is_none() && config.torch_dtype.is_none() {
2165 config.torch_dtype = raw
2166 .get("torch_dtype")
2167 .and_then(serde_json::Value::as_str)
2168 .map(str::to_string);
2169 }
2170 let num_hidden_layers = config.num_hidden_layers.filter(|value| *value > 0)?;
2171 let num_key_value_heads = config.num_key_value_heads.filter(|value| *value > 0)?;
2172 let head_dim = config.head_dim.filter(|value| *value > 0).or_else(|| {
2173 let hidden_size = config.hidden_size?;
2174 let attention_heads = config.num_attention_heads.filter(|value| *value > 0)?;
2175 hidden_size
2176 .is_multiple_of(attention_heads)
2177 .then_some(hidden_size / attention_heads)
2178 })?;
2179 let dtype = config
2180 .kv_cache_dtype
2181 .as_deref()
2182 .or(config.torch_dtype.as_deref());
2183 Some(KvGeometry {
2184 num_hidden_layers,
2185 num_key_value_heads,
2186 head_dim,
2187 kv_bytes_per_element: dtype_bytes(dtype),
2188 hidden_size: config.hidden_size.filter(|value| *value > 0),
2189 })
2190}
2191
2192fn kv_cache_mb_for(geometry: &KvGeometry, context_tokens: usize) -> u64 {
2195 let bytes = 2_u128
2196 .saturating_mul(geometry.num_hidden_layers as u128)
2197 .saturating_mul(geometry.num_key_value_heads as u128)
2198 .saturating_mul(geometry.head_dim as u128)
2199 .saturating_mul(geometry.kv_bytes_per_element as u128)
2200 .saturating_mul(context_tokens as u128);
2201 bytes.div_ceil(BYTES_PER_MB).min(u64::MAX as u128) as u64
2202}
2203
2204fn fallback_kv_cache_mb(model: &ModelSchema, context_tokens: usize) -> u64 {
2209 let active_billions = model_parameter_billions_active(model) as f64;
2210 let bytes_per_token = 147_456.0 * (active_billions / 8.0).max(1.0);
2211 ((bytes_per_token * context_tokens as f64) / BYTES_PER_MB as f64).ceil() as u64
2212}
2213
2214fn prefill_activation_mb(geometry: Option<&KvGeometry>, context_tokens: usize) -> u64 {
2217 let Some(hidden_size) = geometry.and_then(|geometry| geometry.hidden_size) else {
2218 return 0;
2219 };
2220 let bytes = (hidden_size as u128)
2221 .saturating_mul(context_tokens as u128)
2222 .saturating_mul(2)
2223 .saturating_mul(4);
2224 bytes.div_ceil(BYTES_PER_MB).min(u64::MAX as u128) as u64
2225}
2226
2227pub fn estimate_model_memory(
2229 model: &ModelSchema,
2230 hardware: &HardwareInfo,
2231 context_tokens: usize,
2232) -> ModelMemoryEstimate {
2233 estimate_model_memory_with_measured_weights(model, hardware, context_tokens, None)
2234}
2235
2236pub fn estimate_model_memory_for_checkpoint(
2240 model: &ModelSchema,
2241 hardware: &HardwareInfo,
2242 context_tokens: usize,
2243 geometry: &CheckpointGeometry,
2244) -> ModelMemoryEstimate {
2245 estimate_model_memory_with_geometry(model, hardware, context_tokens, None, geometry.get(model))
2246}
2247
2248pub fn estimate_model_memory_with_measured_weights(
2251 model: &ModelSchema,
2252 hardware: &HardwareInfo,
2253 context_tokens: usize,
2254 measured_weights_mb: Option<u64>,
2255) -> ModelMemoryEstimate {
2256 estimate_model_memory_with_geometry(model, hardware, context_tokens, measured_weights_mb, None)
2257}
2258
2259fn estimate_model_memory_with_geometry(
2260 model: &ModelSchema,
2261 hardware: &HardwareInfo,
2262 context_tokens: usize,
2263 measured_weights_mb: Option<u64>,
2264 geometry: Option<&KvGeometry>,
2265) -> ModelMemoryEstimate {
2266 let declared = model.cost.ram_mb.or(model.cost.size_mb);
2267 let (weights_mb, evidence) = if let Some(measured) = measured_weights_mb {
2268 (measured, ModelResourceEvidence::FileSystemMeasured)
2269 } else if let Some(declared) = declared {
2270 (
2271 declared.max(model.cost.size_mb.unwrap_or(0)),
2272 ModelResourceEvidence::CatalogExact,
2273 )
2274 } else {
2275 (
2276 heuristic_weights_mb(model),
2277 ModelResourceEvidence::Heuristic,
2278 )
2279 };
2280 let kv_cache_mb = geometry
2281 .map(|geometry| kv_cache_mb_for(geometry, context_tokens))
2282 .unwrap_or_else(|| fallback_kv_cache_mb(model, context_tokens));
2283 let context_overhead_mb =
2284 kv_cache_mb.saturating_add(prefill_activation_mb(geometry, context_tokens));
2285 let runtime_overhead_mb = backend_runtime_overhead_mb(hardware);
2286 let transient_margin_mb = TRANSIENT_ALLOCATION_MARGIN_MB;
2287 let estimated_peak_mb = weights_mb
2288 .saturating_add(context_overhead_mb)
2289 .saturating_add(runtime_overhead_mb)
2290 .saturating_add(transient_margin_mb);
2291
2292 ModelMemoryEstimate {
2293 weights_mb,
2294 runtime_overhead_mb,
2295 context_overhead_mb,
2296 transient_margin_mb,
2297 estimated_peak_mb,
2298 evidence,
2299 }
2300}
2301
2302fn backend_runtime_overhead_mb(hardware: &HardwareInfo) -> u64 {
2303 match hardware.gpu_backend {
2304 GpuBackend::Metal => METAL_RUNTIME_OVERHEAD_MB,
2305 GpuBackend::Cuda => CUDA_RUNTIME_OVERHEAD_MB,
2306 _ => CPU_RUNTIME_OVERHEAD_MB,
2307 }
2308}
2309
2310fn heuristic_weights_mb(model: &ModelSchema) -> u64 {
2311 let billions = model_parameter_billions_total(model);
2312 (billions as f64 * 600.0).ceil() as u64
2313}
2314
2315pub(crate) fn model_parameter_billions_active(model: &ModelSchema) -> f32 {
2316 model
2317 .param_count
2318 .split_once('(')
2319 .and_then(|(_, rest)| rest.split_once("active"))
2320 .and_then(|(number, _)| parse_parameter_billions(number))
2321 .unwrap_or_else(|| model_parameter_billions_total(model))
2322}
2323
2324pub(crate) fn model_parameter_billions_total(model: &ModelSchema) -> f32 {
2325 parse_parameter_billions(&model.param_count).unwrap_or_else(|| {
2326 let size_mb = model.size_mb();
2327 if size_mb > 0 {
2328 (size_mb as f32 / 600.0).max(0.1)
2329 } else {
2330 0.0
2331 }
2332 })
2333}
2334
2335pub(crate) fn parse_parameter_billions(value: &str) -> Option<f32> {
2336 let value = value.trim();
2337 let number: String = value
2338 .chars()
2339 .take_while(|character| character.is_ascii_digit() || *character == '.')
2340 .collect();
2341 let parsed: f32 = number.parse().ok()?;
2342 if value[number.len()..]
2343 .trim_start()
2344 .to_ascii_lowercase()
2345 .starts_with('m')
2346 {
2347 Some(parsed / 1_000.0)
2348 } else {
2349 Some(parsed)
2350 }
2351}
2352
2353pub fn evaluate_resources(hardware: &HardwareInfo, policy: &ResourcePolicy) -> ResourceEvaluation {
2355 let accelerator_memory = match (hardware.gpu_backend.clone(), hardware.gpu_memory_mb) {
2356 (GpuBackend::Cuda, Some(total_mb)) => Some(AcceleratorResourceBudget {
2357 total_mb,
2358 budget_mb: total_mb,
2359 }),
2360 _ => None,
2361 };
2362
2363 ResourceEvaluation {
2364 host_memory: policy.effective_budget(hardware.total_ram_mb),
2365 accelerator_memory,
2366 }
2367}
2368
2369fn percent_of(total_mb: u64, percent: u64) -> u64 {
2370 let value = (total_mb as u128).saturating_mul(percent as u128) / 100;
2371 value.min(u64::MAX as u128) as u64
2372}
2373
2374fn minimum_emergency_reserve(total_memory_mb: u64) -> u64 {
2375 MINIMUM_EMERGENCY_RESERVE_MB.max(percent_of(total_memory_mb, EMERGENCY_RESERVE_PERCENT))
2376}
2377
2378#[derive(Debug, Error)]
2381pub enum ResourcePolicyError {
2382 #[error("resource policy I/O failed: {0}")]
2383 Io(#[from] io::Error),
2384 #[error("resource policy serialization failed: {0}")]
2385 Serialization(#[from] serde_json::Error),
2386 #[error("Custom model RAM must be finite, nonnegative, and a 0.5 GB increment; got {0}")]
2387 InvalidCustomGigabytes(f64),
2388 #[error("invalid resource policy: {reason}")]
2389 InvalidPolicy { reason: String },
2390}
2391
2392#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
2394#[serde(rename_all = "snake_case")]
2395pub enum ResourcePolicyLoadSource {
2396 Loaded,
2397 MissingDefault,
2398 CorruptDefault,
2399}
2400
2401#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
2403pub struct ResourcePolicyLoadEvidence {
2404 pub policy: ResourcePolicy,
2405 pub source: ResourcePolicyLoadSource,
2406 pub warning: Option<String>,
2407}
2408
2409pub trait ResourcePolicyRepository: Send + Sync {
2410 fn load(&self) -> Result<ResourcePolicy, ResourcePolicyError>;
2411 fn save(&self, policy: &ResourcePolicy) -> Result<(), ResourcePolicyError>;
2412}
2413
2414#[derive(Clone, Debug)]
2416pub struct FileResourcePolicyRepository {
2417 root: PathBuf,
2418}
2419
2420impl Default for FileResourcePolicyRepository {
2421 fn default() -> Self {
2422 Self::new(car_home::root_or_relative())
2423 }
2424}
2425
2426pub const CONCIERGE_POLICY_MARKER: &str = "model-resource-policy.concierge.json";
2429
2430#[derive(Debug, Serialize, Deserialize)]
2431struct ConciergePolicyMarker {
2432 policy: ResourcePolicy,
2433 reason: String,
2434 at: u64,
2435 #[serde(default)]
2437 reverted: bool,
2438}
2439
2440#[derive(Debug, Clone, PartialEq, Eq)]
2442pub struct ConciergeChoice {
2443 pub policy: ResourcePolicy,
2444 pub at: u64,
2445 pub reverted: bool,
2446}
2447
2448impl FileResourcePolicyRepository {
2449 pub fn new(root: PathBuf) -> Self {
2450 Self { root }
2451 }
2452
2453 pub fn concierge_marker(&self) -> Option<ResourcePolicy> {
2458 self.concierge_choice().map(|choice| choice.policy)
2459 }
2460
2461 pub fn concierge_choice(&self) -> Option<ConciergeChoice> {
2464 let bytes = std::fs::read(self.root.join(CONCIERGE_POLICY_MARKER)).ok()?;
2465 serde_json::from_slice::<ConciergePolicyMarker>(&bytes)
2466 .ok()
2467 .map(|marker| ConciergeChoice {
2468 policy: marker.policy,
2469 at: marker.at,
2470 reverted: marker.reverted,
2471 })
2472 }
2473
2474 pub fn save_concierge_choice(
2476 &self,
2477 policy: &ResourcePolicy,
2478 reason: &str,
2479 at: u64,
2480 reverted: bool,
2481 ) -> Result<(), ResourcePolicyError> {
2482 ResourcePolicyRepository::save(self, policy)?;
2483 let marker = ConciergePolicyMarker {
2484 policy: policy.clone(),
2485 reason: reason.to_string(),
2486 at,
2487 reverted,
2488 };
2489 let path = self.root.join(CONCIERGE_POLICY_MARKER);
2490 let tmp = self.root.join(format!(".{CONCIERGE_POLICY_MARKER}.tmp"));
2491 let bytes = serde_json::to_vec_pretty(&marker).map_err(io::Error::other)?;
2492 std::fs::write(&tmp, bytes)?;
2493 std::fs::rename(&tmp, &path)?;
2494 Ok(())
2495 }
2496
2497 pub fn clear_concierge_marker(&self) -> Result<(), ResourcePolicyError> {
2499 match std::fs::remove_file(self.root.join(CONCIERGE_POLICY_MARKER)) {
2500 Err(error) if error.kind() != io::ErrorKind::NotFound => Err(error.into()),
2501 _ => Ok(()),
2502 }
2503 }
2504
2505 pub fn path(&self) -> PathBuf {
2506 self.root.join(RESOURCE_POLICY_FILE)
2507 }
2508
2509 pub fn load_with_evidence(&self) -> Result<ResourcePolicyLoadEvidence, ResourcePolicyError> {
2512 let path = self.path();
2513 let file = match open_resource_policy(&path) {
2514 Ok(file) => file,
2515 Err(error) if error.kind() == io::ErrorKind::NotFound => {
2516 return Ok(ResourcePolicyLoadEvidence {
2517 policy: ResourcePolicy::everyday(),
2518 source: ResourcePolicyLoadSource::MissingDefault,
2519 warning: None,
2520 });
2521 }
2522 Err(error) => {
2523 if std::fs::symlink_metadata(&path)
2524 .is_ok_and(|metadata| !metadata.is_file() || metadata.file_type().is_symlink())
2525 {
2526 return Ok(corrupt_default(
2527 "The saved resource policy could not be loaded because it is not a regular file.",
2528 ));
2529 }
2530 return Err(error.into());
2531 }
2532 };
2533 let metadata = file.metadata()?;
2534 if !metadata.is_file() {
2535 return Ok(corrupt_default(
2536 "The saved resource policy could not be loaded because it is not a regular file.",
2537 ));
2538 }
2539 if metadata.len() > MAX_POLICY_BYTES {
2540 return Ok(corrupt_default(
2541 "The saved resource policy could not be loaded because it exceeds the size limit.",
2542 ));
2543 }
2544
2545 let mut raw = Vec::new();
2546 file.take(MAX_POLICY_BYTES + 1).read_to_end(&mut raw)?;
2547 if raw.len() as u64 > MAX_POLICY_BYTES {
2548 return Ok(corrupt_default(
2549 "The saved resource policy could not be loaded because it exceeds the size limit.",
2550 ));
2551 }
2552 let policy = match serde_json::from_slice::<ResourcePolicy>(&raw) {
2553 Ok(policy) => policy,
2554 Err(error) => {
2555 return Ok(corrupt_default(format!(
2556 "The saved resource policy could not be loaded: {error}"
2557 )));
2558 }
2559 };
2560 if let Err(error) = policy.validate() {
2561 return Ok(corrupt_default(format!(
2562 "The saved resource policy could not be loaded: {error}"
2563 )));
2564 }
2565
2566 Ok(ResourcePolicyLoadEvidence {
2567 policy,
2568 source: ResourcePolicyLoadSource::Loaded,
2569 warning: None,
2570 })
2571 }
2572}
2573
2574fn open_resource_policy(path: &Path) -> io::Result<std::fs::File> {
2575 let mut options = std::fs::OpenOptions::new();
2576 options.read(true);
2577 #[cfg(unix)]
2578 {
2579 use std::os::unix::fs::OpenOptionsExt;
2580 options.custom_flags(libc::O_NOFOLLOW);
2581 }
2582 options.open(path)
2583}
2584
2585fn corrupt_default(warning: impl Into<String>) -> ResourcePolicyLoadEvidence {
2586 ResourcePolicyLoadEvidence {
2587 policy: ResourcePolicy::everyday(),
2588 source: ResourcePolicyLoadSource::CorruptDefault,
2589 warning: Some(warning.into()),
2590 }
2591}
2592
2593impl ResourcePolicyRepository for FileResourcePolicyRepository {
2594 fn load(&self) -> Result<ResourcePolicy, ResourcePolicyError> {
2595 Ok(self.load_with_evidence()?.policy)
2596 }
2597
2598 fn save(&self, policy: &ResourcePolicy) -> Result<(), ResourcePolicyError> {
2599 policy.validate()?;
2600 let _guard = mutation_lock()
2601 .lock()
2602 .unwrap_or_else(std::sync::PoisonError::into_inner);
2603 ensure_private_directory(&self.root)?;
2604
2605 let temp_path = unique_temp_path(&self.root);
2606 let result = (|| {
2607 let body = serde_json::to_vec_pretty(policy)?;
2608 let mut file = open_private_temp(&temp_path)?;
2609 file.write_all(&body)?;
2610 file.sync_all()?;
2611 atomic_replace(&temp_path, &self.path())?;
2612 car_secrets::harden_owner_only(&self.path());
2613 sync_directory(&self.root)?;
2614 Ok(())
2615 })();
2616 if result.is_err() {
2617 let _ = std::fs::remove_file(&temp_path);
2618 }
2619 result
2620 }
2621}
2622
2623fn mutation_lock() -> &'static Mutex<()> {
2624 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
2625 LOCK.get_or_init(|| Mutex::new(()))
2626}
2627
2628fn unique_temp_path(root: &Path) -> PathBuf {
2629 static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(0);
2630 let sequence = NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed);
2631 let epoch_nanos = std::time::SystemTime::now()
2632 .duration_since(std::time::UNIX_EPOCH)
2633 .map(|duration| duration.as_nanos())
2634 .unwrap_or(0);
2635 root.join(format!(
2636 ".{RESOURCE_POLICY_FILE}.{}.{}.{}.tmp",
2637 std::process::id(),
2638 epoch_nanos,
2639 sequence
2640 ))
2641}
2642
2643fn ensure_private_directory(path: &Path) -> io::Result<()> {
2644 std::fs::create_dir_all(path)?;
2645 #[cfg(unix)]
2646 {
2647 use std::os::unix::fs::PermissionsExt;
2648 std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
2649 }
2650 car_secrets::harden_owner_only(path);
2651 Ok(())
2652}
2653
2654#[cfg(unix)]
2655fn open_private_temp(path: &Path) -> io::Result<std::fs::File> {
2656 use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
2657
2658 let file = std::fs::OpenOptions::new()
2659 .write(true)
2660 .create_new(true)
2661 .mode(0o600)
2662 .open(path)?;
2663 file.set_permissions(std::fs::Permissions::from_mode(0o600))?;
2664 Ok(file)
2665}
2666
2667#[cfg(not(unix))]
2668fn open_private_temp(path: &Path) -> io::Result<std::fs::File> {
2669 let file = std::fs::OpenOptions::new()
2670 .write(true)
2671 .create_new(true)
2672 .open(path)?;
2673 car_secrets::harden_owner_only(path);
2674 Ok(file)
2675}
2676
2677#[cfg(not(windows))]
2678fn atomic_replace(source: &Path, destination: &Path) -> io::Result<()> {
2679 std::fs::rename(source, destination)
2680}
2681
2682#[cfg(windows)]
2683fn atomic_replace(source: &Path, destination: &Path) -> io::Result<()> {
2684 use std::os::windows::ffi::OsStrExt;
2685
2686 const MOVEFILE_REPLACE_EXISTING: u32 = 0x1;
2687 const MOVEFILE_WRITE_THROUGH: u32 = 0x8;
2688 #[link(name = "kernel32")]
2689 unsafe extern "system" {
2690 fn MoveFileExW(existing: *const u16, new: *const u16, flags: u32) -> i32;
2691 }
2692 let source = source
2693 .as_os_str()
2694 .encode_wide()
2695 .chain(Some(0))
2696 .collect::<Vec<_>>();
2697 let destination = destination
2698 .as_os_str()
2699 .encode_wide()
2700 .chain(Some(0))
2701 .collect::<Vec<_>>();
2702 let replaced = unsafe {
2705 MoveFileExW(
2706 source.as_ptr(),
2707 destination.as_ptr(),
2708 MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
2709 )
2710 };
2711 if replaced == 0 {
2712 Err(io::Error::last_os_error())
2713 } else {
2714 Ok(())
2715 }
2716}
2717
2718#[cfg(unix)]
2719fn sync_directory(path: &Path) -> io::Result<()> {
2720 std::fs::File::open(path)?.sync_all()
2721}
2722
2723#[cfg(not(unix))]
2724fn sync_directory(_path: &Path) -> io::Result<()> {
2725 Ok(())
2726}
2727
2728#[cfg(test)]
2729mod tests {
2730 use super::*;
2731
2732 fn hardware(
2733 total_ram_mb: u64,
2734 gpu_backend: crate::hardware::GpuBackend,
2735 vram_mb: Option<u64>,
2736 ) -> crate::hardware::HardwareInfo {
2737 crate::hardware::HardwareInfo {
2738 os: "test".into(),
2739 arch: "test".into(),
2740 cpu_cores: 8,
2741 total_ram_mb,
2742 gpu_backend,
2743 gpu_memory_mb: vram_mb,
2744 gpu_devices: Vec::new(),
2745 recommended_model: "fixture".into(),
2746 recommended_context: 4_096,
2747 max_model_mb: total_ram_mb,
2748 }
2749 }
2750
2751 #[test]
2752 fn profiles_compute_exact_32_gb_budgets() {
2753 let total = 32 * 1024;
2754
2755 assert_eq!(
2756 ResourcePolicy::everyday()
2757 .effective_budget(total)
2758 .configured_model_ceiling_mb,
2759 13_107
2760 );
2761 assert_eq!(
2762 ResourcePolicy::local_focused()
2763 .effective_budget(total)
2764 .configured_model_ceiling_mb,
2765 26_214
2766 );
2767 assert_eq!(
2768 ResourcePolicy::custom_gb(12.5)
2769 .unwrap()
2770 .effective_budget(total)
2771 .configured_model_ceiling_mb,
2772 12_800
2773 );
2774 assert_eq!(
2775 ResourcePolicy::everyday().recommendation_target_mb(total),
2776 6_553
2777 );
2778 }
2779
2780 #[test]
2781 fn model_memory_estimate_keeps_transient_margin_distinct_and_totals_exactly() {
2782 let catalog = crate::registry::builtin_catalog();
2783 let model = catalog
2784 .iter()
2785 .find(|model| model.id == "mlx/qwen3-4b:4bit")
2786 .unwrap();
2787 let estimate = estimate_model_memory(
2788 model,
2789 &hardware(32 * 1024, GpuBackend::Metal, None),
2790 RECOMMENDATION_CONTEXT_TOKENS,
2791 );
2792
2793 assert_eq!(estimate.evidence, ModelResourceEvidence::CatalogExact);
2794 assert_eq!(estimate.weights_mb, 2_400);
2795 assert_eq!(estimate.runtime_overhead_mb, 512);
2796 assert_eq!(estimate.transient_margin_mb, 1_024);
2797 assert_eq!(
2798 estimate.estimated_peak_mb,
2799 estimate.weights_mb
2800 + estimate.context_overhead_mb
2801 + estimate.runtime_overhead_mb
2802 + estimate.transient_margin_mb
2803 );
2804
2805 let measured = estimate_model_memory_with_measured_weights(
2806 model,
2807 &hardware(32 * 1024, GpuBackend::Metal, None),
2808 RECOMMENDATION_CONTEXT_TOKENS,
2809 Some(2_321),
2810 );
2811 assert_eq!(measured.weights_mb, 2_321);
2812 assert_eq!(measured.evidence, ModelResourceEvidence::FileSystemMeasured);
2813 }
2814
2815 fn qwen3_8b_model() -> ModelSchema {
2816 crate::registry::builtin_catalog()
2817 .into_iter()
2818 .find(|model| model.id == "mlx/qwen3-8b:4bit")
2819 .unwrap()
2820 }
2821
2822 fn qwen3_8b_geometry() -> KvGeometry {
2823 KvGeometry {
2824 num_hidden_layers: 36,
2825 num_key_value_heads: 8,
2826 head_dim: 128,
2827 kv_bytes_per_element: 2,
2828 hidden_size: Some(4_096),
2829 }
2830 }
2831
2832 #[test]
2833 fn kv_cache_qwen3_8b_geometry_matches_binary_mb_table() {
2834 let geometry = qwen3_8b_geometry();
2835 assert_eq!(kv_cache_mb_for(&geometry, 9_000), 1_266);
2836 assert_eq!(kv_cache_mb_for(&geometry, 13_096), 1_842);
2837 assert_eq!(prefill_activation_mb(Some(&geometry), 9_000), 282);
2838 }
2839
2840 #[test]
2841 fn kv_cache_unknown_geometry_uses_conservative_qwen3_8b_floor() {
2842 let model = qwen3_8b_model();
2843 let estimate =
2844 estimate_model_memory(&model, &hardware(32 * 1024, GpuBackend::Metal, None), 9_000);
2845
2846 assert_eq!(estimate.context_overhead_mb, 1_266);
2847 }
2848
2849 #[test]
2850 fn kv_cache_unknown_geometry_scales_models_larger_than_8b() {
2851 let mut model = qwen3_8b_model();
2852 model.param_count = "32B".into();
2853
2854 let estimate =
2855 estimate_model_memory(&model, &hardware(64 * 1024, GpuBackend::Metal, None), 9_000);
2856
2857 assert_eq!(estimate.context_overhead_mb, 5_063);
2858 }
2859
2860 #[test]
2861 fn kv_geometry_installed_config_drives_admission_context_charge() {
2862 let root = tempfile::tempdir().unwrap();
2863 let model = qwen3_8b_model();
2864 let model_dir = root.path().join(&model.name);
2865 std::fs::create_dir_all(&model_dir).unwrap();
2866 std::fs::write(
2867 model_dir.join("config.json"),
2868 r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32,"torch_dtype":"bfloat16"}"#,
2869 )
2870 .unwrap();
2871 let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2872 ResourcePolicy::local_focused(),
2873 hardware(32 * 1024, GpuBackend::Metal, None),
2874 Arc::new(FixedLiveMemoryProbe::known(24_000)),
2875 Some(root.path().to_path_buf()),
2876 );
2877
2878 let estimate = coordinator.preflight(&model, 9_000).estimate;
2879
2880 assert_eq!(estimate.context_overhead_mb, 1_548);
2881 }
2882
2883 #[test]
2888 fn recommendation_estimate_uses_the_installed_geometry_admission_uses() {
2889 let root = tempfile::tempdir().unwrap();
2890 let model = qwen3_8b_model();
2891 let model_dir = root.path().join(&model.name);
2892 std::fs::create_dir_all(&model_dir).unwrap();
2893 std::fs::write(
2894 model_dir.join("config.json"),
2895 r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32,"torch_dtype":"bfloat16"}"#,
2896 )
2897 .unwrap();
2898 let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
2899 let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2900 ResourcePolicy::local_focused(),
2901 hardware.clone(),
2902 Arc::new(FixedLiveMemoryProbe::known(24_000)),
2903 Some(root.path().to_path_buf()),
2904 );
2905 let admitted = coordinator
2906 .preflight(&model, RECOMMENDATION_CONTEXT_TOKENS)
2907 .estimate
2908 .estimated_peak_mb;
2909 let peak = |geometry: &CheckpointGeometry| {
2910 estimate_model_memory_for_checkpoint(
2911 &model,
2912 &hardware,
2913 RECOMMENDATION_CONTEXT_TOKENS,
2914 geometry,
2915 )
2916 .estimated_peak_mb
2917 };
2918
2919 assert_eq!(
2920 peak(&CheckpointGeometry::read([&model], root.path())),
2921 admitted
2922 );
2923 assert_eq!(peak(&coordinator.checkpoint_geometry([&model])), admitted);
2924 assert_ne!(peak(&CheckpointGeometry::default()), admitted);
2926 assert_eq!(
2927 peak(&CheckpointGeometry::default()),
2928 estimate_model_memory(&model, &hardware, RECOMMENDATION_CONTEXT_TOKENS)
2929 .estimated_peak_mb
2930 );
2931 }
2932
2933 #[test]
2934 fn kv_geometry_reads_a_multimodal_checkpoints_nested_text_config() {
2935 let root = tempfile::tempdir().unwrap();
2936 let model = qwen3_8b_model();
2937 let model_dir = root.path().join(&model.name);
2938 std::fs::create_dir_all(&model_dir).unwrap();
2939 std::fs::write(
2941 model_dir.join("config.json"),
2942 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}}"#,
2943 )
2944 .unwrap();
2945 let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2946 ResourcePolicy::local_focused(),
2947 hardware(32 * 1024, GpuBackend::Metal, None),
2948 Arc::new(FixedLiveMemoryProbe::known(24_000)),
2949 Some(root.path().to_path_buf()),
2950 );
2951 assert_eq!(
2952 coordinator
2953 .preflight(&model, 9_000)
2954 .estimate
2955 .context_overhead_mb,
2956 1_548
2957 );
2958 }
2959
2960 #[test]
2961 fn kv_geometry_cache_forgets_missing_config_after_install() {
2962 let root = tempfile::tempdir().unwrap();
2963 let model = qwen3_8b_model();
2964 let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2965 ResourcePolicy::local_focused(),
2966 hardware(32 * 1024, GpuBackend::Metal, None),
2967 Arc::new(FixedLiveMemoryProbe::known(24_000)),
2968 Some(root.path().to_path_buf()),
2969 );
2970
2971 let before_install = coordinator.preflight(&model, 9_000).estimate;
2972 let model_dir = root.path().join(&model.name);
2973 std::fs::create_dir_all(&model_dir).unwrap();
2974 std::fs::write(
2975 model_dir.join("config.json"),
2976 r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096,"num_attention_heads":32,"torch_dtype":"bfloat16"}"#,
2977 )
2978 .unwrap();
2979 let after_install = coordinator.preflight(&model, 9_000).estimate;
2980
2981 assert_eq!(before_install.context_overhead_mb, 1_266);
2982 assert_eq!(after_install.context_overhead_mb, 1_548);
2983 }
2984
2985 #[test]
2986 fn kv_geometry_derives_head_dim_and_memoizes_installed_config() {
2987 let root = tempfile::tempdir().unwrap();
2988 let model = qwen3_8b_model();
2989 let model_dir = root.path().join(&model.name);
2990 std::fs::create_dir_all(&model_dir).unwrap();
2991 let config_path = model_dir.join("config.json");
2992 std::fs::write(
2993 &config_path,
2994 r#"{"num_hidden_layers":36,"num_key_value_heads":8,"hidden_size":4096,"num_attention_heads":32}"#,
2995 )
2996 .unwrap();
2997 let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
2998 ResourcePolicy::local_focused(),
2999 hardware(32 * 1024, GpuBackend::Metal, None),
3000 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3001 Some(root.path().to_path_buf()),
3002 );
3003
3004 let first = coordinator.preflight(&model, 9_000).estimate;
3005 let modified = std::fs::metadata(&config_path).unwrap().modified().unwrap();
3008 let len = std::fs::metadata(&config_path).unwrap().len() as usize;
3009 std::fs::write(&config_path, "x".repeat(len)).unwrap();
3010 std::fs::File::options()
3011 .write(true)
3012 .open(&config_path)
3013 .unwrap()
3014 .set_modified(modified)
3015 .unwrap();
3016 let second = coordinator.preflight(&model, 9_000).estimate;
3017
3018 assert_eq!(first.context_overhead_mb, 1_548);
3019 assert_eq!(second.context_overhead_mb, first.context_overhead_mb);
3020 }
3021
3022 #[test]
3023 fn kv_geometry_cache_is_scoped_to_coordinator_models_dir() {
3024 let first_root = tempfile::tempdir().unwrap();
3025 let second_root = tempfile::tempdir().unwrap();
3026 let model = qwen3_8b_model();
3027 for (root, config) in [
3028 (
3029 first_root.path(),
3030 r#"{"num_hidden_layers":36,"num_key_value_heads":8,"head_dim":128,"hidden_size":4096}"#,
3031 ),
3032 (
3033 second_root.path(),
3034 r#"{"num_hidden_layers":2,"num_key_value_heads":2,"head_dim":64,"hidden_size":256}"#,
3035 ),
3036 ] {
3037 let model_dir = root.join(&model.name);
3038 std::fs::create_dir_all(&model_dir).unwrap();
3039 std::fs::write(model_dir.join("config.json"), config).unwrap();
3040 }
3041 let coordinator = |models_dir: &Path| {
3042 LocalAdmissionCoordinator::with_probe_and_models_dir(
3043 ResourcePolicy::local_focused(),
3044 hardware(32 * 1024, GpuBackend::Metal, None),
3045 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3046 Some(models_dir.to_path_buf()),
3047 )
3048 };
3049
3050 assert_eq!(
3054 coordinator(first_root.path())
3055 .preflight(&model, 9_000)
3056 .estimate
3057 .context_overhead_mb,
3058 1_548
3059 );
3060 assert_eq!(
3061 coordinator(second_root.path())
3062 .preflight(&model, 9_000)
3063 .estimate
3064 .context_overhead_mb,
3065 27
3066 );
3067 }
3068
3069 #[test]
3070 fn kv_geometry_malformed_config_falls_back_without_panicking() {
3071 let root = tempfile::tempdir().unwrap();
3072 let model = qwen3_8b_model();
3073 let model_dir = root.path().join(&model.name);
3074 std::fs::create_dir_all(&model_dir).unwrap();
3075 std::fs::write(model_dir.join("config.json"), b"not json").unwrap();
3076 let coordinator = LocalAdmissionCoordinator::with_probe_and_models_dir(
3077 ResourcePolicy::local_focused(),
3078 hardware(32 * 1024, GpuBackend::Metal, None),
3079 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3080 Some(root.path().to_path_buf()),
3081 );
3082
3083 let estimate = coordinator.preflight(&model, 9_000).estimate;
3084
3085 assert_eq!(estimate.context_overhead_mb, 1_266);
3086 }
3087
3088 #[test]
3089 fn kv_cache_estimate_is_monotone_and_keeps_non_context_terms_fixed() {
3090 let model = qwen3_8b_model();
3091 let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
3092 let short = estimate_model_memory(&model, &hardware, 9_000);
3093 let long = estimate_model_memory(&model, &hardware, 13_096);
3094
3095 assert!(long.context_overhead_mb >= short.context_overhead_mb);
3096 assert_eq!(short.weights_mb, long.weights_mb);
3097 assert_eq!(short.runtime_overhead_mb, long.runtime_overhead_mb);
3098 assert_eq!(short.transient_margin_mb, long.transient_margin_mb);
3099 }
3100
3101 #[test]
3102 fn custom_zero_disables_loads_and_overlarge_values_clamp_below_emergency_reserve() {
3103 let total = 32 * 1024;
3104
3105 assert_eq!(
3106 ResourcePolicy::custom_gb(0.0)
3107 .unwrap()
3108 .effective_budget(total)
3109 .effective_new_load_ceiling_mb,
3110 0
3111 );
3112 let result = ResourcePolicy::custom_gb(99.0)
3113 .unwrap()
3114 .effective_budget(total);
3115 assert_eq!(result.emergency_reserve_mb, 3_276);
3116 assert_eq!(result.configured_model_ceiling_mb, total - 3_276);
3117 assert!(result.normalization_notice.is_some());
3118 }
3119
3120 #[test]
3121 fn custom_gigabytes_reject_non_half_steps_negative_and_non_finite_values() {
3122 for invalid in [10.3, -0.5, f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
3123 assert!(
3124 ResourcePolicy::custom_gb(invalid).is_err(),
3125 "accepted {invalid:?}"
3126 );
3127 }
3128
3129 assert_eq!(
3130 ResourcePolicy::custom_gb(10.5).unwrap().custom_max_model_mb,
3131 Some(10_752)
3132 );
3133 }
3134
3135 #[test]
3136 fn repository_round_trips_exact_half_gb_and_uses_private_atomic_files() {
3137 let dir = tempfile::tempdir().unwrap();
3138 let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3139
3140 repository
3141 .save(&ResourcePolicy::custom_gb(10.5).unwrap())
3142 .unwrap();
3143
3144 assert_eq!(
3145 repository.path(),
3146 dir.path().join("model-resource-policy.json")
3147 );
3148 assert_eq!(repository.load().unwrap().custom_max_model_mb, Some(10_752));
3149 assert_private_mode(&repository.path(), 0o600);
3150 assert_private_mode(dir.path(), 0o700);
3151 assert_no_temp_files(dir.path());
3152 }
3153
3154 #[test]
3155 fn repository_missing_or_corrupt_file_falls_back_without_deleting_source() {
3156 let dir = tempfile::tempdir().unwrap();
3157 let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3158 assert_eq!(repository.load().unwrap(), ResourcePolicy::everyday());
3159
3160 let corrupt = br#"{"profile":"custom","custom_max_model_mb":"broken"}"#;
3161 std::fs::write(repository.path(), corrupt).unwrap();
3162
3163 assert_eq!(repository.load().unwrap(), ResourcePolicy::everyday());
3164 assert_eq!(std::fs::read(repository.path()).unwrap(), corrupt);
3165 }
3166
3167 #[test]
3168 fn repository_rejects_invalid_policy_shapes_before_writing() {
3169 let dir = tempfile::tempdir().unwrap();
3170 let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3171 let invalid = [
3172 ResourcePolicy {
3173 profile: ResourceProfile::Custom,
3174 custom_max_model_mb: None,
3175 },
3176 ResourcePolicy {
3177 profile: ResourceProfile::Custom,
3178 custom_max_model_mb: Some(1),
3179 },
3180 ResourcePolicy {
3181 profile: ResourceProfile::Custom,
3182 custom_max_model_mb: Some(513),
3183 },
3184 ResourcePolicy {
3185 profile: ResourceProfile::Everyday,
3186 custom_max_model_mb: Some(512),
3187 },
3188 ResourcePolicy {
3189 profile: ResourceProfile::LocalFocused,
3190 custom_max_model_mb: Some(512),
3191 },
3192 ];
3193
3194 for policy in invalid {
3195 let error = repository.save(&policy).unwrap_err();
3196 assert!(matches!(error, ResourcePolicyError::InvalidPolicy { .. }));
3197 assert!(!repository.path().exists());
3198 }
3199 }
3200
3201 #[test]
3202 fn load_evidence_distinguishes_loaded_missing_and_corrupt_defaults() {
3203 let dir = tempfile::tempdir().unwrap();
3204 let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3205
3206 let missing = repository.load_with_evidence().unwrap();
3207 assert_eq!(missing.policy, ResourcePolicy::everyday());
3208 assert_eq!(missing.source, ResourcePolicyLoadSource::MissingDefault);
3209 assert!(missing.warning.is_none());
3210
3211 repository.save(&ResourcePolicy::local_focused()).unwrap();
3212 let loaded = repository.load_with_evidence().unwrap();
3213 assert_eq!(loaded.policy, ResourcePolicy::local_focused());
3214 assert_eq!(loaded.source, ResourcePolicyLoadSource::Loaded);
3215 assert!(loaded.warning.is_none());
3216
3217 let corrupt = br#"{"profile":"custom","custom_max_model_mb":"broken"}"#;
3218 std::fs::write(repository.path(), corrupt).unwrap();
3219 let recovered = repository.load_with_evidence().unwrap();
3220 assert_eq!(recovered.policy, ResourcePolicy::everyday());
3221 assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
3222 assert!(recovered
3223 .warning
3224 .as_deref()
3225 .is_some_and(|warning| { warning.contains("could not be loaded") }));
3226 assert_eq!(std::fs::read(repository.path()).unwrap(), corrupt);
3227 }
3228
3229 #[test]
3230 fn corrupt_evidence_covers_unknown_fields_invalid_shapes_and_oversized_files() {
3231 let dir = tempfile::tempdir().unwrap();
3232 let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3233 let corrupt_documents = [
3234 br#"{"#.to_vec(),
3235 br#"{"profile":"everyday","custom_max_model_mb":null,"extra":true}"#.to_vec(),
3236 br#"{"profile":"custom","custom_max_model_mb":null}"#.to_vec(),
3237 br#"{"profile":"custom","custom_max_model_mb":1}"#.to_vec(),
3238 br#"{"profile":"custom","custom_max_model_mb":513}"#.to_vec(),
3239 br#"{"profile":"everyday","custom_max_model_mb":512}"#.to_vec(),
3240 br#"{"profile":"local_focused","custom_max_model_mb":512}"#.to_vec(),
3241 vec![b' '; MAX_POLICY_BYTES as usize + 1],
3242 ];
3243
3244 for document in corrupt_documents {
3245 std::fs::write(repository.path(), &document).unwrap();
3246 let recovered = repository.load_with_evidence().unwrap();
3247 assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
3248 assert!(recovered.warning.is_some());
3249 assert_eq!(std::fs::read(repository.path()).unwrap(), document);
3250 }
3251 }
3252
3253 #[cfg(unix)]
3254 #[test]
3255 fn non_regular_policy_source_is_reported_as_corrupt_without_following_it() {
3256 use std::os::unix::fs::symlink;
3257
3258 let dir = tempfile::tempdir().unwrap();
3259 let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
3260 let target = dir.path().join("target.json");
3261 std::fs::write(
3262 &target,
3263 br#"{"profile":"local_focused","custom_max_model_mb":null}"#,
3264 )
3265 .unwrap();
3266 symlink(&target, repository.path()).unwrap();
3267
3268 let recovered = repository.load_with_evidence().unwrap();
3269 assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
3270 assert!(recovered.warning.is_some());
3271 assert!(repository.path().is_symlink());
3272 }
3273
3274 #[test]
3275 fn concurrent_saves_leave_one_complete_document_and_no_staging_files() {
3276 let dir = tempfile::tempdir().unwrap();
3277 let repository =
3278 std::sync::Arc::new(FileResourcePolicyRepository::new(dir.path().to_path_buf()));
3279 let mut writers = Vec::new();
3280 for index in 0..16_u64 {
3281 let repository = repository.clone();
3282 writers.push(std::thread::spawn(move || {
3283 repository
3284 .save(&ResourcePolicy {
3285 profile: ResourceProfile::Custom,
3286 custom_max_model_mb: Some(index * 512),
3287 })
3288 .unwrap();
3289 }));
3290 }
3291 for writer in writers {
3292 writer.join().unwrap();
3293 }
3294
3295 let loaded = repository.load().unwrap();
3296 assert_eq!(loaded.profile, ResourceProfile::Custom);
3297 assert!(loaded.custom_max_model_mb.unwrap().is_multiple_of(512));
3298 assert_no_temp_files(dir.path());
3299 }
3300
3301 #[test]
3302 fn cuda_uses_separate_vram_fit_and_host_ram_policy() {
3303 let hardware = hardware(
3304 64 * 1024,
3305 crate::hardware::GpuBackend::Cuda,
3306 Some(12 * 1024),
3307 );
3308
3309 let evidence = evaluate_resources(&hardware, &ResourcePolicy::everyday());
3310
3311 assert_eq!(evidence.host_memory.configured_model_ceiling_mb, 26_214);
3312 let accelerator = evidence.accelerator_memory.unwrap();
3313 assert_eq!(accelerator.total_mb, 12 * 1024);
3314 assert_eq!(accelerator.budget_mb, 12 * 1024);
3315 assert_ne!(
3316 accelerator.budget_mb,
3317 evidence.host_memory.configured_model_ceiling_mb
3318 );
3319 }
3320
3321 #[test]
3322 fn resource_preflight_zero_custom_budget_blocks_new_load() {
3323 let coordinator = LocalAdmissionCoordinator::with_probe(
3324 ResourcePolicy::custom_gb(0.0).unwrap(),
3325 hardware(32 * 1024, GpuBackend::Metal, None),
3326 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3327 );
3328 let model = crate::registry::builtin_catalog()
3329 .into_iter()
3330 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3331 .unwrap();
3332
3333 let preflight = coordinator.preflight(&model, 2_048);
3334 assert_eq!(preflight.verdict, LocalLoadVerdict::DisabledByPolicy);
3335 }
3336
3337 #[test]
3338 fn resource_preflight_zero_budget_does_not_kill_resident_inference() {
3339 let coordinator = LocalAdmissionCoordinator::with_probe(
3340 ResourcePolicy::custom_gb(0.0).unwrap(),
3341 hardware(32 * 1024, GpuBackend::Metal, None),
3342 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3343 );
3344 let model = crate::registry::builtin_catalog()
3345 .into_iter()
3346 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3347 .unwrap();
3348 coordinator.mark_resident(&model.id, 2_400);
3349
3350 let preflight = coordinator.preflight(&model, 2_048);
3351 assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
3352 assert_eq!(
3353 preflight.estimated_incremental_mb,
3354 preflight.estimate.context_overhead_mb
3355 + preflight.estimate.runtime_overhead_mb
3356 + preflight.estimate.transient_margin_mb
3357 );
3358 }
3359
3360 #[test]
3361 fn resource_preflight_resident_weights_are_incremental_only() {
3362 let coordinator = LocalAdmissionCoordinator::with_probe(
3363 ResourcePolicy::custom_gb(8.0).unwrap(),
3364 hardware(32 * 1024, GpuBackend::Metal, None),
3365 Arc::new(FixedLiveMemoryProbe::known(6_000)),
3366 );
3367 let model = crate::registry::builtin_catalog()
3368 .into_iter()
3369 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3370 .unwrap();
3371 coordinator.mark_resident(&model.id, 2_400);
3372
3373 let preflight = coordinator.preflight(&model, 2_048);
3374 assert_eq!(
3375 preflight.estimated_incremental_mb,
3376 preflight.estimate.context_overhead_mb
3377 + preflight.estimate.runtime_overhead_mb
3378 + preflight.estimate.transient_margin_mb
3379 );
3380 assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
3381 }
3382
3383 #[test]
3384 fn resource_preflight_unavailable_live_probe_is_explicit() {
3385 let coordinator = LocalAdmissionCoordinator::with_probe(
3386 ResourcePolicy::everyday(),
3387 hardware(32 * 1024, GpuBackend::Metal, None),
3388 Arc::new(FixedLiveMemoryProbe::unknown()),
3389 );
3390 let model = crate::registry::builtin_catalog()
3391 .into_iter()
3392 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3393 .unwrap();
3394
3395 let preflight = coordinator.preflight(&model, 2_048);
3396 assert_eq!(preflight.live_available_mb, None);
3397 assert_eq!(preflight.verdict, LocalLoadVerdict::LiveMemoryUnknown);
3398 }
3399
3400 #[test]
3401 fn resource_preflight_simultaneous_reservations_are_atomic() {
3402 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3403 ResourcePolicy::custom_gb(5.0).unwrap(),
3404 hardware(32 * 1024, GpuBackend::Metal, None),
3405 Arc::new(FixedLiveMemoryProbe::known(6_000)),
3406 ));
3407 let model = crate::registry::builtin_catalog()
3408 .into_iter()
3409 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3410 .unwrap();
3411 coordinator.mark_resident(&model.id, 2_400);
3412
3413 let first = coordinator.reserve(&model, 2_048).unwrap();
3414 let second = coordinator.reserve(&model, 2_048).unwrap_err();
3415 assert_eq!(
3416 second.preflight.verdict,
3417 LocalLoadVerdict::ExceedsConfiguredCeiling
3418 );
3419 drop(first);
3420 assert!(coordinator.reserve(&model, 2_048).is_ok());
3421 }
3422
3423 #[test]
3424 fn distinct_state_roots_reserve_atomically_against_one_machine_ledger() {
3425 let machine_ledger = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
3426 let policy = ResourcePolicy::custom_gb(6.0).unwrap();
3427 let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
3428 let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3429 policy.clone(),
3430 hardware.clone(),
3431 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3432 machine_ledger.clone(),
3433 ));
3434 let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3435 policy,
3436 hardware,
3437 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3438 machine_ledger,
3439 ));
3440 let mut model = crate::registry::builtin_catalog()
3441 .into_iter()
3442 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3443 .unwrap();
3444 model.cost.ram_mb = Some(2 * 1024);
3445 model.cost.size_mb = Some(2 * 1024);
3446 let barrier = Arc::new(std::sync::Barrier::new(2));
3447
3448 let attempts = [first, second].map(|coordinator| {
3449 let model = model.clone();
3450 let barrier = barrier.clone();
3451 std::thread::spawn(move || {
3452 barrier.wait();
3453 coordinator.reserve(&model, 0)
3454 })
3455 });
3456 let outcomes = attempts.map(|attempt| attempt.join().unwrap());
3457
3458 assert_eq!(outcomes.iter().filter(|result| result.is_ok()).count(), 1);
3459 assert_eq!(
3460 outcomes
3461 .iter()
3462 .find_map(|result| result.as_ref().err())
3463 .expect("one cross-root request must be blocked")
3464 .preflight
3465 .verdict,
3466 LocalLoadVerdict::ExceedsConfiguredCeiling
3467 );
3468 }
3469
3470 #[test]
3471 fn resource_preflight_cuda_charges_weights_to_vram_and_overhead_to_host() {
3472 let coordinator = LocalAdmissionCoordinator::with_probe(
3473 ResourcePolicy::custom_gb(2.0).unwrap(),
3474 hardware(64 * 1024, GpuBackend::Cuda, Some(12 * 1024)),
3475 Arc::new(FixedLiveMemoryProbe::known(20_000)),
3476 );
3477 let model = crate::registry::builtin_catalog()
3478 .into_iter()
3479 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3480 .unwrap();
3481
3482 let preflight = coordinator.preflight(&model, 2_048);
3483 assert_eq!(preflight.resident_model_mb, 0);
3484 assert_eq!(preflight.accelerator_total_mb, Some(12 * 1024));
3485 assert_eq!(
3486 preflight.accelerator_incremental_mb,
3487 Some(preflight.estimate.weights_mb)
3488 );
3489 assert_eq!(
3490 preflight.estimated_incremental_mb,
3491 preflight.estimate.context_overhead_mb
3492 + preflight.estimate.runtime_overhead_mb
3493 + preflight.estimate.transient_margin_mb
3494 );
3495 assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
3496 }
3497
3498 #[test]
3499 fn resource_preflight_cuda_vram_reservations_are_atomic() {
3500 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3501 ResourcePolicy::custom_gb(8.0).unwrap(),
3502 hardware(64 * 1024, GpuBackend::Cuda, Some(3_000)),
3503 Arc::new(FixedLiveMemoryProbe::known(20_000)),
3504 ));
3505 let model = crate::registry::builtin_catalog()
3506 .into_iter()
3507 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3508 .unwrap();
3509
3510 let first = coordinator.reserve(&model, 2_048).unwrap();
3511 assert_eq!(
3512 coordinator
3513 .reserve(&model, 2_048)
3514 .unwrap_err()
3515 .preflight
3516 .verdict,
3517 LocalLoadVerdict::ExceedsConfiguredCeiling
3518 );
3519 drop(first);
3520 assert!(coordinator.reserve(&model, 2_048).is_ok());
3521 }
3522
3523 #[test]
3524 fn resource_preflight_model_maintenance_races_atomically_with_reserve() {
3525 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3526 ResourcePolicy::everyday(),
3527 hardware(32 * 1024, GpuBackend::Metal, None),
3528 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3529 ));
3530 let model = crate::registry::builtin_catalog()
3531 .into_iter()
3532 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3533 .unwrap();
3534
3535 let active = coordinator.reserve(&model, 2_048).unwrap();
3536 assert!(matches!(
3537 coordinator.begin_model_maintenance(&model.id),
3538 Err(ModelMaintenanceError::ModelInUse(_))
3539 ));
3540 drop(active);
3541
3542 let maintenance = coordinator.begin_model_maintenance(&model.id).unwrap();
3543 assert_eq!(
3544 coordinator
3545 .reserve(&model, 2_048)
3546 .unwrap_err()
3547 .preflight
3548 .verdict,
3549 LocalLoadVerdict::ModelMaintenance
3550 );
3551 drop(maintenance);
3552 assert!(coordinator.reserve(&model, 2_048).is_ok());
3553 }
3554
3555 #[test]
3556 fn measured_weights_are_rechecked_atomically_before_allocation() {
3557 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3558 ResourcePolicy::custom_gb(5.0).unwrap(),
3559 hardware(32 * 1024, GpuBackend::Metal, None),
3560 Arc::new(FixedLiveMemoryProbe::known(16_000)),
3561 ));
3562 let model = crate::registry::builtin_catalog()
3563 .into_iter()
3564 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3565 .unwrap();
3566 let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
3567
3568 let blocked = reservation
3569 .reconcile_measured_weights(6 * 1024 * 1024 * 1024)
3570 .unwrap_err();
3571
3572 assert_eq!(
3573 blocked.preflight.verdict,
3574 LocalLoadVerdict::ExceedsConfiguredCeiling
3575 );
3576 assert_eq!(
3577 coordinator.preflight(&model, 2_048).active_reservations_mb,
3578 reservation.reserved_incremental_mb(),
3579 "a rejected resize must not mutate the live reservation"
3580 );
3581 }
3582
3583 #[test]
3584 fn cache_publication_transfers_cold_reservation_without_double_counting() {
3585 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3586 ResourcePolicy::custom_gb(8.0).unwrap(),
3587 hardware(32 * 1024, GpuBackend::Metal, None),
3588 Arc::new(FixedLiveMemoryProbe::known(16_000)),
3589 ));
3590 let model = crate::registry::builtin_catalog()
3591 .into_iter()
3592 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3593 .unwrap();
3594 let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
3595 let measured_bytes = 3 * 1024 * 1024 * 1024_u64;
3596
3597 reservation
3598 .reconcile_measured_weights(measured_bytes)
3599 .unwrap();
3600 reservation.publish_resident_weights(measured_bytes);
3601 let after = coordinator.preflight(&model, 2_048);
3602 let request_overhead = after.estimate.context_overhead_mb
3603 + after.estimate.runtime_overhead_mb
3604 + after.estimate.transient_margin_mb;
3605
3606 assert_eq!(after.resident_model_mb, 3 * 1024);
3607 assert_eq!(after.active_reservations_mb, request_overhead);
3608 assert_eq!(reservation.reserved_incremental_mb(), request_overhead);
3609 }
3610
3611 #[test]
3612 fn resident_publication_never_reprobes_or_rejects_after_allocation() {
3613 struct CountingProbe {
3614 calls: AtomicU64,
3615 available_mb: u64,
3616 }
3617
3618 impl LiveMemoryProbe for CountingProbe {
3619 fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
3620 self.calls.fetch_add(1, Ordering::Relaxed);
3621 Ok(Some(self.available_mb))
3622 }
3623 }
3624
3625 let probe = Arc::new(CountingProbe {
3626 calls: AtomicU64::new(0),
3627 available_mb: 24_000,
3628 });
3629 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3630 ResourcePolicy::custom_gb(8.0).unwrap(),
3631 hardware(32 * 1024, GpuBackend::Metal, None),
3632 probe.clone(),
3633 ));
3634 let model = crate::registry::builtin_catalog()
3635 .into_iter()
3636 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3637 .unwrap();
3638 let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
3639 let measured = 3 * 1024 * 1024 * 1024_u64;
3640 reservation.reconcile_measured_weights(measured).unwrap();
3641 let calls_before_publish = probe.calls.load(Ordering::Relaxed);
3642
3643 reservation.publish_resident_weights(measured + 512 * 1024 * 1024);
3644
3645 assert_eq!(probe.calls.load(Ordering::Relaxed), calls_before_publish);
3646 assert!(coordinator.is_resident(&model.id));
3647 assert_eq!(
3648 coordinator
3649 .state
3650 .lock()
3651 .unwrap_or_else(std::sync::PoisonError::into_inner)
3652 .resident_models
3653 .get(&model.id)
3654 .map(|allocation| allocation.weights_mb),
3655 Some(3 * 1024 + 512)
3656 );
3657 }
3658
3659 #[test]
3660 fn simultaneous_cold_reservations_converge_on_one_resident_allocation() {
3661 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3662 ResourcePolicy::custom_gb(16.0).unwrap(),
3663 hardware(32 * 1024, GpuBackend::Metal, None),
3664 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3665 ));
3666 let model = crate::registry::builtin_catalog()
3667 .into_iter()
3668 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3669 .unwrap();
3670 let mut first = coordinator.reserve(&model, 2_048).unwrap();
3671 let mut second = coordinator.reserve(&model, 2_048).unwrap();
3672 let measured = 3 * 1024 * 1024 * 1024_u64;
3673
3674 first.reconcile_measured_weights(measured).unwrap();
3675 second.reconcile_measured_weights(measured).unwrap();
3676 first.publish_resident_weights(measured);
3677 second.publish_resident_weights(measured);
3678
3679 let after = coordinator.preflight(&model, 2_048);
3680 let per_request_overhead = after.estimate.context_overhead_mb
3681 + after.estimate.runtime_overhead_mb
3682 + after.estimate.transient_margin_mb;
3683 assert_eq!(after.resident_model_mb, 3 * 1024);
3684 assert_eq!(
3685 after.active_reservations_mb,
3686 per_request_overhead * 2,
3687 "peer publication must release only this reservation's redundant cold weights"
3688 );
3689 }
3690
3691 #[test]
3692 fn resident_eviction_before_allocation_repromotes_reservation_to_cold() {
3693 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3694 ResourcePolicy::custom_gb(8.0).unwrap(),
3695 hardware(32 * 1024, GpuBackend::Cpu, None),
3696 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3697 ));
3698 let model = crate::registry::builtin_catalog()
3699 .into_iter()
3700 .find(|model| model.id == "mlx/qwen3-4b:4bit")
3701 .unwrap();
3702 let measured = 3 * 1024 * 1024 * 1024_u64;
3703 let mut first = coordinator.reserve(&model, 2_048).unwrap();
3704 first.reconcile_measured_weights(measured).unwrap();
3705 first.publish_resident_weights(measured);
3706 drop(first);
3707 let mut replacement = coordinator.reserve(&model, 2_048).unwrap();
3708 coordinator.mark_evicted(&model.id);
3709
3710 let preflight = replacement.reconcile_measured_weights(measured).unwrap();
3711 assert!(preflight.estimated_incremental_mb >= 3 * 1024);
3712 }
3713
3714 #[test]
3715 fn maintenance_catalog_alias_blocks_provider_alias_reservation() {
3716 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3717 ResourcePolicy::custom_gb(8.0).unwrap(),
3718 hardware(32 * 1024, GpuBackend::Cpu, None),
3719 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3720 ));
3721 coordinator
3722 .register_model_aliases("mlx/kokoro-82m:6bit", ["mlx-community/Kokoro-82M-6bit"]);
3723 let _maintenance = coordinator
3724 .begin_model_maintenance("mlx/kokoro-82m:6bit")
3725 .unwrap();
3726 let mut provider_schema = crate::registry::builtin_catalog()
3727 .into_iter()
3728 .find(|model| model.is_local())
3729 .expect("one local schema");
3730 provider_schema.id = "mlx-community/Kokoro-82M-6bit".into();
3731
3732 let blocked = coordinator.reserve(&provider_schema, 512).unwrap_err();
3733 assert_eq!(
3734 blocked.preflight.verdict,
3735 LocalLoadVerdict::ModelMaintenance
3736 );
3737 assert_eq!(blocked.preflight.model_id, "mlx/kokoro-82m:6bit");
3738 }
3739
3740 #[test]
3741 fn late_alias_registration_rekeys_active_and_resident_state_atomically() {
3742 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3743 ResourcePolicy::custom_gb(8.0).unwrap(),
3744 hardware(32 * 1024, GpuBackend::Cpu, None),
3745 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3746 ));
3747 let mut schema = crate::registry::builtin_catalog()
3748 .into_iter()
3749 .find(|model| model.is_local())
3750 .unwrap();
3751 schema.id = "provider/artifact".into();
3752 let active = coordinator.reserve(&schema, 0).unwrap();
3753
3754 coordinator.register_model_aliases("catalog/model:default", ["provider/artifact"]);
3755 assert!(matches!(
3756 coordinator.begin_model_maintenance("catalog/model:default"),
3757 Err(ModelMaintenanceError::ModelInUse(_))
3758 ));
3759 drop(active);
3760
3761 let maintenance = coordinator
3762 .begin_model_maintenance("catalog/model:default")
3763 .unwrap();
3764 let error = coordinator.reserve(&schema, 0).unwrap_err();
3765 assert_eq!(error.preflight.verdict, LocalLoadVerdict::ModelMaintenance);
3766 drop(maintenance);
3767 }
3768
3769 #[test]
3770 fn process_allocations_teardown_by_exact_owner_without_sibling_erasure() {
3771 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3772 ResourcePolicy::custom_gb(16.0).unwrap(),
3773 hardware(32 * 1024, GpuBackend::Cpu, None),
3774 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3775 ));
3776 let logical = "local/shared-model";
3777 let worker = worker_process_allocation_id(logical);
3778 let vllm = vllm_process_allocation_id(logical);
3779 let bytes = 512 * 1024 * 1024_u64;
3780 let mut worker_load = coordinator
3781 .reserve_measured_host_allocation(logical, &worker, bytes, 0)
3782 .unwrap();
3783 worker_load.publish_resident_weights_as(&worker, bytes);
3784 drop(worker_load);
3785 let mut vllm_load = coordinator
3786 .reserve_measured_host_allocation(logical, &vllm, bytes, 0)
3787 .unwrap();
3788 vllm_load.publish_resident_weights_as(&vllm, bytes);
3789 drop(vllm_load);
3790 assert_eq!(coordinator.resident_model_mb(), 1024);
3791
3792 coordinator.mark_teardown_pending_allocation(logical, &worker);
3793 coordinator.mark_teardown_pending_allocation(logical, &vllm);
3794 coordinator.finish_teardown_allocation(logical, &worker);
3795 assert!(coordinator.teardown_pending(logical));
3796 assert_eq!(
3797 coordinator.resident_allocation_ids(logical),
3798 vec![vllm.clone()]
3799 );
3800 assert_eq!(coordinator.resident_model_mb(), 512);
3801
3802 coordinator.finish_teardown_allocation(logical, &vllm);
3803 assert!(!coordinator.teardown_pending(logical));
3804 assert!(!coordinator.is_resident(logical));
3805 }
3806
3807 #[test]
3808 fn replacement_generation_reconcile_removes_peer_resident_discount_before_load() {
3809 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3810 ResourcePolicy::custom_gb(1.0).unwrap(),
3811 hardware(32 * 1024, GpuBackend::Cpu, None),
3812 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3813 ));
3814 let logical = "local/replaced-worker";
3815 let bytes = 512 * 1024 * 1024_u64;
3816 let mut old = coordinator
3817 .reserve_measured_host_allocation(logical, "worker:old", bytes, 0)
3818 .unwrap();
3819 old.publish_resident_weights_as("worker:old", bytes);
3820 drop(old);
3821
3822 let mut replacement = coordinator
3823 .reserve_measured_host(logical, bytes, 0)
3824 .expect("logical preflight initially sees the old resident");
3825 assert_eq!(replacement.reserved_incremental_mb(), 0);
3826 replacement.bind_allocation_id("worker:new");
3827 replacement
3828 .reconcile_measured_weights(bytes)
3829 .expect("the exact replacement generation fits by itself");
3830 assert_eq!(replacement.reserved_incremental_mb(), 512);
3831 assert!(
3832 coordinator
3833 .reserve_measured_host("different/model", bytes, 0)
3834 .is_err(),
3835 "a second cold model must see both old residency and the replacement generation"
3836 );
3837
3838 drop(replacement);
3839 assert!(coordinator
3840 .reserve_measured_host("different/model", bytes, 0)
3841 .is_ok());
3842 }
3843
3844 #[test]
3845 fn pre_ack_pending_charge_blocks_other_scope_until_exact_exit_ack() {
3846 let machine = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
3847 let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3848 ResourcePolicy::custom_gb(2.0).unwrap(),
3849 hardware(32 * 1024, GpuBackend::Cpu, None),
3850 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3851 machine.clone(),
3852 ));
3853 let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3854 ResourcePolicy::custom_gb(2.0).unwrap(),
3855 hardware(32 * 1024, GpuBackend::Cpu, None),
3856 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3857 machine,
3858 ));
3859 let logical = "managed/starting";
3860 let allocation = worker_process_allocation_id(logical);
3861 let cold = first
3862 .reserve_measured_host_allocation(logical, &allocation, 1024 * 1024 * 1024, 0)
3863 .unwrap();
3864 first.mark_teardown_pending_allocation_with_charge(
3865 logical,
3866 &allocation,
3867 cold.reconciled_weights_bytes(),
3868 );
3869 drop(cold);
3870
3871 let blocked = second
3872 .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3873 .unwrap_err();
3874 assert_eq!(
3875 blocked.preflight.verdict,
3876 LocalLoadVerdict::ExceedsConfiguredCeiling
3877 );
3878 first.finish_teardown_allocation(logical, "worker:unrelated-sibling");
3879 assert!(first.teardown_pending(logical));
3880 assert_eq!(first.resident_model_mb(), 1024);
3881
3882 first.finish_teardown_allocation(logical, &allocation);
3883 assert!(!first.teardown_pending(logical));
3884 assert!(second
3885 .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3886 .is_ok());
3887 }
3888
3889 #[test]
3890 fn detached_native_lease_keeps_machine_charge_after_request_cancellation() {
3891 let machine = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
3892 let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3893 ResourcePolicy::custom_gb(2.0).unwrap(),
3894 hardware(32 * 1024, GpuBackend::Cpu, None),
3895 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3896 machine.clone(),
3897 ));
3898 let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
3899 ResourcePolicy::custom_gb(2.0).unwrap(),
3900 hardware(32 * 1024, GpuBackend::Cpu, None),
3901 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3902 machine,
3903 ));
3904 let reservation = first
3905 .reserve_measured_host("detached/model", 1024 * 1024 * 1024, 0)
3906 .unwrap();
3907 let detached = reservation.detached_lease();
3908 drop(reservation);
3909 assert!(second
3910 .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3911 .is_err());
3912 drop(detached);
3913 assert!(second
3914 .reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
3915 .is_ok());
3916 }
3917
3918 #[test]
3919 fn normal_awaited_detached_work_shares_one_charge_with_its_request() {
3920 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3921 ResourcePolicy::custom_gb(1.0).unwrap(),
3922 hardware(32 * 1024, GpuBackend::Cpu, None),
3923 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3924 ));
3925 let request = coordinator
3926 .reserve_measured_host("native/model-a", 512 * 1024 * 1024, 0)
3927 .unwrap();
3928 let detached = request.detached_lease();
3929
3930 let peer = coordinator
3931 .reserve_measured_host("native/model-b", 512 * 1024 * 1024, 0)
3932 .expect("request + its detached job are one 512 MB allocation, not two");
3933 assert_eq!(
3934 coordinator
3935 .state
3936 .lock()
3937 .unwrap_or_else(std::sync::PoisonError::into_inner)
3938 .active_host_reservations_mb,
3939 1024
3940 );
3941
3942 drop(peer);
3943 drop(detached);
3944 assert_eq!(
3945 coordinator
3946 .state
3947 .lock()
3948 .unwrap_or_else(std::sync::PoisonError::into_inner)
3949 .active_host_reservations_mb,
3950 512
3951 );
3952 drop(request);
3953 assert_eq!(
3954 coordinator
3955 .state
3956 .lock()
3957 .unwrap_or_else(std::sync::PoisonError::into_inner)
3958 .active_host_reservations_mb,
3959 0
3960 );
3961 }
3962
3963 #[test]
3964 fn concurrent_non_aligned_starts_transfer_cold_weights_to_pending_without_double_charge() {
3965 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
3966 ResourcePolicy::custom_gb(8.0).unwrap(),
3967 hardware(32 * 1024, GpuBackend::Cpu, None),
3968 Arc::new(FixedLiveMemoryProbe::known(24_000)),
3969 ));
3970 let first_allocation = vllm_process_allocation_id("managed/model-a");
3971 let second_allocation = vllm_process_allocation_id("managed/model-b");
3972 let mib = 1024 * 1024_u64;
3973 let first_measured_bytes = 1024 * mib + 1;
3974 let second_measured_bytes = 512 * mib + 1;
3975 let mut first = coordinator
3976 .reserve_measured_host_allocation(
3977 "managed/model-a",
3978 &first_allocation,
3979 first_measured_bytes,
3980 128,
3981 )
3982 .unwrap();
3983 let mut second = coordinator
3984 .reserve_measured_host_allocation(
3985 "managed/model-b",
3986 &second_allocation,
3987 second_measured_bytes,
3988 256,
3989 )
3990 .unwrap();
3991
3992 first.transfer_cold_weights_to_pending_allocation(&first_allocation, first_measured_bytes);
3993 second
3994 .transfer_cold_weights_to_pending_allocation(&second_allocation, second_measured_bytes);
3995
3996 assert_eq!(
3997 coordinator.resident_model_mb(),
3998 1538,
3999 "each non-MiB-aligned allocation must be rounded up independently"
4000 );
4001 assert_eq!(
4002 coordinator
4003 .state
4004 .lock()
4005 .unwrap_or_else(std::sync::PoisonError::into_inner)
4006 .active_host_reservations_mb,
4007 384,
4008 "pending weights replace cold request weights while request overhead remains active"
4009 );
4010
4011 first.publish_resident_weights_as(&first_allocation, first_measured_bytes);
4012 second.publish_resident_weights_as(&second_allocation, second_measured_bytes);
4013 assert_eq!(coordinator.resident_model_mb(), 1538);
4014 assert_eq!(
4015 coordinator
4016 .state
4017 .lock()
4018 .unwrap_or_else(std::sync::PoisonError::into_inner)
4019 .active_host_reservations_mb,
4020 384,
4021 "publication must not subtract transferred weights twice"
4022 );
4023
4024 drop(first);
4025 drop(second);
4026 assert_eq!(
4027 coordinator
4028 .state
4029 .lock()
4030 .unwrap_or_else(std::sync::PoisonError::into_inner)
4031 .active_host_reservations_mb,
4032 0
4033 );
4034 assert_eq!(coordinator.resident_model_mb(), 1538);
4035 }
4036
4037 #[test]
4038 fn pending_allocation_identity_is_idempotent_and_sibling_exact() {
4039 let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
4040 ResourcePolicy::custom_gb(8.0).unwrap(),
4041 hardware(32 * 1024, GpuBackend::Cpu, None),
4042 Arc::new(FixedLiveMemoryProbe::known(24_000)),
4043 ));
4044 coordinator.mark_teardown_pending_allocation_with_charge(
4045 "same/model",
4046 "worker:same/model",
4047 512 * 1024 * 1024,
4048 );
4049 coordinator.mark_teardown_pending_allocation_with_charge(
4050 "same/model",
4051 "worker:same/model",
4052 512 * 1024 * 1024,
4053 );
4054 coordinator.mark_teardown_pending_allocation_with_charge(
4055 "same/model",
4056 "vllm:same/model",
4057 256 * 1024 * 1024,
4058 );
4059 assert_eq!(coordinator.resident_model_mb(), 768);
4060 coordinator.finish_teardown_allocation("same/model", "worker:same/model");
4061 assert!(coordinator.teardown_pending("same/model"));
4062 assert_eq!(coordinator.resident_model_mb(), 256);
4063 coordinator.finish_teardown_allocation("same/model", "worker:same/model");
4064 assert_eq!(coordinator.resident_model_mb(), 256);
4065 coordinator.finish_teardown_allocation("same/model", "vllm:same/model");
4066 assert!(!coordinator.teardown_pending("same/model"));
4067 }
4068
4069 #[test]
4070 fn scoped_coordinator_identity_is_stable_for_one_state_root() {
4071 let root = tempfile::tempdir().unwrap();
4072 let first = scoped_local_admission(
4073 root.path(),
4074 ResourcePolicy::custom_gb(4.0).unwrap(),
4075 hardware(32 * 1024, GpuBackend::Cpu, None),
4076 );
4077 let second = scoped_local_admission(
4078 root.path(),
4079 ResourcePolicy::custom_gb(8.0).unwrap(),
4080 hardware(32 * 1024, GpuBackend::Cpu, None),
4081 );
4082 assert!(Arc::ptr_eq(&first, &second));
4083 assert_eq!(second.policy(), ResourcePolicy::custom_gb(8.0).unwrap());
4084 }
4085
4086 #[cfg(unix)]
4087 #[test]
4088 fn scoped_coordinator_identity_unifies_symlinked_state_roots() {
4089 use std::os::unix::fs::symlink;
4090
4091 let fixture = tempfile::tempdir().unwrap();
4092 let real = fixture.path().join("real-state");
4093 std::fs::create_dir(&real).unwrap();
4094 let alias = fixture.path().join("state-alias");
4095 symlink(&real, &alias).unwrap();
4096
4097 let first = scoped_local_admission(
4098 &real,
4099 ResourcePolicy::custom_gb(4.0).unwrap(),
4100 hardware(32 * 1024, GpuBackend::Cpu, None),
4101 );
4102 let second = scoped_local_admission(
4103 &alias,
4104 ResourcePolicy::custom_gb(8.0).unwrap(),
4105 hardware(32 * 1024, GpuBackend::Cpu, None),
4106 );
4107 assert!(Arc::ptr_eq(&first, &second));
4108 assert_eq!(second.policy(), ResourcePolicy::custom_gb(8.0).unwrap());
4109 }
4110
4111 #[test]
4112 fn unavailable_state_root_normalization_is_absolute_and_lexically_stable() {
4113 let fixture = tempfile::tempdir().unwrap();
4114 let missing = fixture.path().join("not-created").join("..").join("state");
4115 assert_eq!(
4116 normalized_state_root_key(&missing),
4117 normalized_state_root_key(&fixture.path().join("state"))
4118 );
4119 }
4120
4121 #[cfg(unix)]
4122 #[test]
4123 fn missing_leaf_under_symlinked_parent_keeps_one_scope_identity() {
4124 use std::os::unix::fs::symlink;
4125
4126 let fixture = tempfile::tempdir().unwrap();
4127 let real = fixture.path().join("real");
4128 std::fs::create_dir(&real).unwrap();
4129 let alias = fixture.path().join("alias");
4130 symlink(&real, &alias).unwrap();
4131
4132 assert_eq!(
4133 normalized_state_root_key(&alias.join("missing").join("state")),
4134 normalized_state_root_key(&real.join("missing").join("state"))
4135 );
4136 }
4137
4138 #[cfg(unix)]
4139 #[test]
4140 fn parent_components_are_resolved_after_symlinks_not_lexically_before_them() {
4141 use std::os::unix::fs::symlink;
4142
4143 let fixture = tempfile::tempdir().unwrap();
4144 let physical_parent = fixture.path().join("physical");
4145 let physical_child = physical_parent.join("child");
4146 std::fs::create_dir_all(&physical_child).unwrap();
4147 let aliases = fixture.path().join("aliases");
4148 std::fs::create_dir(&aliases).unwrap();
4149 let alias = aliases.join("runtime");
4150 symlink(&physical_child, &alias).unwrap();
4151
4152 let through_alias = alias.join("..").join("missing-state");
4153 assert_eq!(
4154 normalized_state_root_key(&through_alias),
4155 normalized_state_root_key(&physical_parent.join("missing-state"))
4156 );
4157 assert_ne!(
4158 normalized_state_root_key(&through_alias),
4159 aliases.join("missing-state")
4160 );
4161 }
4162
4163 fn assert_no_temp_files(directory: &std::path::Path) {
4164 let entries = std::fs::read_dir(directory)
4165 .unwrap()
4166 .map(|entry| entry.unwrap().path())
4167 .collect::<Vec<_>>();
4168 assert_eq!(entries, vec![directory.join("model-resource-policy.json")]);
4169 }
4170
4171 #[cfg(unix)]
4172 fn assert_private_mode(path: &std::path::Path, expected: u32) {
4173 use std::os::unix::fs::PermissionsExt;
4174 assert_eq!(
4175 std::fs::metadata(path).unwrap().permissions().mode() & 0o777,
4176 expected
4177 );
4178 }
4179
4180 #[cfg(not(unix))]
4181 fn assert_private_mode(_path: &std::path::Path, _expected: u32) {}
4182}