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