use std::collections::{HashMap, HashSet};
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Weak};
use std::sync::{Mutex, OnceLock};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use crate::hardware::{GpuBackend, HardwareInfo};
use crate::schema::ModelSchema;
pub const RESOURCE_POLICY_FILE: &str = "model-resource-policy.json";
const EVERYDAY_MODEL_PERCENT: u64 = 40;
const LOCAL_FOCUSED_MODEL_PERCENT: u64 = 80;
const EMERGENCY_RESERVE_PERCENT: u64 = 10;
const MINIMUM_EMERGENCY_RESERVE_MB: u64 = 2 * 1024;
const MAX_POLICY_BYTES: u64 = 64 * 1024;
pub const RECOMMENDATION_CONTEXT_TOKENS: usize = 8_192;
const METAL_RUNTIME_OVERHEAD_MB: u64 = 512;
const CUDA_RUNTIME_OVERHEAD_MB: u64 = 512;
const CPU_RUNTIME_OVERHEAD_MB: u64 = 1_024;
const TRANSIENT_ALLOCATION_MARGIN_MB: u64 = 1_024;
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ResourceProfile {
Everyday,
LocalFocused,
Custom,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct ResourcePolicy {
pub profile: ResourceProfile,
pub custom_max_model_mb: Option<u64>,
}
impl Default for ResourcePolicy {
fn default() -> Self {
Self::everyday()
}
}
impl ResourcePolicy {
pub fn everyday() -> Self {
Self {
profile: ResourceProfile::Everyday,
custom_max_model_mb: None,
}
}
pub fn local_focused() -> Self {
Self {
profile: ResourceProfile::LocalFocused,
custom_max_model_mb: None,
}
}
pub fn custom_gb(gigabytes: f64) -> Result<Self, ResourcePolicyError> {
if !gigabytes.is_finite() || gigabytes < 0.0 {
return Err(ResourcePolicyError::InvalidCustomGigabytes(gigabytes));
}
let half_gb_steps = gigabytes * 2.0;
if half_gb_steps.fract() != 0.0 || half_gb_steps > u64::MAX as f64 {
return Err(ResourcePolicyError::InvalidCustomGigabytes(gigabytes));
}
let steps = half_gb_steps as u64;
let custom_max_model_mb = steps
.checked_mul(512)
.ok_or(ResourcePolicyError::InvalidCustomGigabytes(gigabytes))?;
Ok(Self {
profile: ResourceProfile::Custom,
custom_max_model_mb: Some(custom_max_model_mb),
})
}
pub fn effective_budget(&self, total_memory_mb: u64) -> EffectiveResourceBudget {
let emergency_reserve_mb = minimum_emergency_reserve(total_memory_mb);
let safe_maximum_mb = total_memory_mb.saturating_sub(emergency_reserve_mb);
let requested_ceiling_mb = match self.profile {
ResourceProfile::Everyday => percent_of(total_memory_mb, EVERYDAY_MODEL_PERCENT),
ResourceProfile::LocalFocused => {
percent_of(total_memory_mb, LOCAL_FOCUSED_MODEL_PERCENT)
}
ResourceProfile::Custom => self.custom_max_model_mb.unwrap_or(0),
};
let configured_model_ceiling_mb = requested_ceiling_mb.min(safe_maximum_mb);
let normalization_notice = (matches!(self.profile, ResourceProfile::Custom)
&& requested_ceiling_mb > safe_maximum_mb)
.then(|| {
format!(
"The saved Custom allocation was adjusted from {requested_ceiling_mb} MB to \
{safe_maximum_mb} MB on this machine to preserve the \
{emergency_reserve_mb} MB emergency reserve."
)
});
EffectiveResourceBudget {
total_memory_mb,
emergency_reserve_mb,
configured_model_ceiling_mb,
effective_new_load_ceiling_mb: configured_model_ceiling_mb,
normalization_notice,
}
}
pub fn recommendation_target_mb(&self, total_memory_mb: u64) -> u64 {
let ceiling = self
.effective_budget(total_memory_mb)
.configured_model_ceiling_mb;
if self.profile == ResourceProfile::Everyday {
ceiling / 2
} else {
ceiling
}
}
pub fn validate(&self) -> Result<(), ResourcePolicyError> {
match self.profile {
ResourceProfile::Custom => match self.custom_max_model_mb {
Some(value) if value.is_multiple_of(512) => Ok(()),
Some(value) => Err(ResourcePolicyError::InvalidPolicy {
reason: format!(
"Custom model RAM must be a 0.5 GB (512 MB) increment; got {value} MB"
),
}),
None => Err(ResourcePolicyError::InvalidPolicy {
reason: "Custom profile requires custom_max_model_mb".into(),
}),
},
ResourceProfile::Everyday | ResourceProfile::LocalFocused => {
if self.custom_max_model_mb.is_none() {
Ok(())
} else {
Err(ResourcePolicyError::InvalidPolicy {
reason: format!(
"{:?} profile must not set custom_max_model_mb",
self.profile
),
})
}
}
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EffectiveResourceBudget {
pub total_memory_mb: u64,
pub emergency_reserve_mb: u64,
pub configured_model_ceiling_mb: u64,
pub effective_new_load_ceiling_mb: u64,
pub normalization_notice: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AcceleratorResourceBudget {
pub total_mb: u64,
pub budget_mb: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResourceEvaluation {
pub host_memory: EffectiveResourceBudget,
pub accelerator_memory: Option<AcceleratorResourceBudget>,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ModelResourceEvidence {
CatalogExact,
FileSystemMeasured,
Heuristic,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct ModelMemoryEstimate {
pub weights_mb: u64,
pub runtime_overhead_mb: u64,
pub context_overhead_mb: u64,
pub transient_margin_mb: u64,
pub estimated_peak_mb: u64,
pub evidence: ModelResourceEvidence,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LocalLoadPreflight {
pub model_id: String,
pub estimate: ModelMemoryEstimate,
pub configured_ceiling_mb: u64,
pub resident_model_mb: u64,
pub active_reservations_mb: u64,
pub estimated_incremental_mb: u64,
pub accelerator_total_mb: Option<u64>,
pub accelerator_resident_mb: Option<u64>,
pub accelerator_incremental_mb: Option<u64>,
pub live_available_mb: Option<u64>,
pub emergency_reserve_mb: u64,
pub verdict: LocalLoadVerdict,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum LocalLoadVerdict {
Allowed,
LiveMemoryUnknown,
DisabledByPolicy,
ExceedsConfiguredCeiling,
InsufficientLiveMemory,
ModelMaintenance,
PendingTeardown,
}
impl LocalLoadVerdict {
pub fn permits_static_fallback(&self) -> bool {
matches!(self, Self::Allowed | Self::LiveMemoryUnknown)
}
}
pub trait LiveMemoryProbe: Send + Sync {
fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError>;
}
#[derive(Default)]
pub struct SystemLiveMemoryProbe;
impl LiveMemoryProbe for SystemLiveMemoryProbe {
fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
Ok(crate::hardware::available_ram_mb())
}
}
#[cfg(test)]
struct FixedLiveMemoryProbe(Option<u64>);
#[cfg(test)]
impl FixedLiveMemoryProbe {
fn known(available_mb: u64) -> Self {
Self(Some(available_mb))
}
fn unknown() -> Self {
Self(None)
}
}
#[cfg(test)]
impl LiveMemoryProbe for FixedLiveMemoryProbe {
fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
Ok(self.0)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum WeightPlacement {
Host,
Accelerator,
}
#[derive(Clone, Debug)]
struct ResidentAllocation {
weights_mb: u64,
placement: WeightPlacement,
logical_model_id: String,
}
#[derive(Clone, Default)]
struct AdmissionState {
resident_models: HashMap<String, ResidentAllocation>,
active_host_reservations_mb: u64,
active_accelerator_reservations_mb: u64,
active_by_model: HashMap<String, usize>,
maintenance_models: HashSet<String>,
pending_teardown_models: HashMap<String, HashSet<String>>,
model_aliases: HashMap<String, HashSet<String>>,
next_request_id: u64,
}
#[derive(Clone, Default)]
struct MachineAdmissionLedger {
resident_models: HashMap<(u64, String), ResidentAllocation>,
pending_allocations: HashMap<(u64, String), ResidentAllocation>,
active_host_by_owner: HashMap<u64, u64>,
active_accelerator_by_owner: HashMap<u64, u64>,
}
fn process_machine_admission_ledger() -> Arc<Mutex<MachineAdmissionLedger>> {
static LEDGER: OnceLock<Arc<Mutex<MachineAdmissionLedger>>> = OnceLock::new();
LEDGER
.get_or_init(|| Arc::new(Mutex::new(MachineAdmissionLedger::default())))
.clone()
}
fn next_admission_owner_id() -> u64 {
static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
}
pub fn worker_process_allocation_id(model_id: &str) -> String {
format!("worker:{model_id}")
}
pub fn vllm_process_allocation_id(model_id: &str) -> String {
format!("vllm:{model_id}")
}
pub struct LocalAdmissionCoordinator {
policy: std::sync::RwLock<ResourcePolicy>,
hardware: HardwareInfo,
live_probe: Arc<dyn LiveMemoryProbe>,
state: Mutex<AdmissionState>,
machine_ledger: Arc<Mutex<MachineAdmissionLedger>>,
resident_activity_leases: Mutex<HashMap<String, Arc<crate::model_management::ModelLease>>>,
owner_id: u64,
}
impl Drop for LocalAdmissionCoordinator {
fn drop(&mut self) {
let mut machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
machine
.resident_models
.retain(|(owner_id, _), _| *owner_id != self.owner_id);
machine
.pending_allocations
.retain(|(owner_id, _), _| *owner_id != self.owner_id);
machine.active_host_by_owner.remove(&self.owner_id);
machine.active_accelerator_by_owner.remove(&self.owner_id);
}
}
fn scoped_admission_registry() -> &'static Mutex<HashMap<PathBuf, Weak<LocalAdmissionCoordinator>>>
{
static REGISTRY: OnceLock<Mutex<HashMap<PathBuf, Weak<LocalAdmissionCoordinator>>>> =
OnceLock::new();
REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}
pub fn normalized_state_root_key(state_root: &Path) -> PathBuf {
let absolute = if state_root.is_absolute() {
state_root.to_path_buf()
} else if let Ok(current) = std::env::current_dir() {
current.join(state_root)
} else {
return state_root.to_path_buf();
};
if let Ok(canonical) = std::fs::canonicalize(&absolute) {
return canonical;
}
let mut ancestor = absolute.clone();
let mut missing_suffix = Vec::new();
while let Some(component) = ancestor.components().next_back() {
let name = match component {
std::path::Component::Normal(name) => name.to_owned(),
std::path::Component::CurDir => std::ffi::OsString::from("."),
std::path::Component::ParentDir => std::ffi::OsString::from(".."),
std::path::Component::RootDir | std::path::Component::Prefix(_) => break,
};
missing_suffix.push(name);
ancestor.pop();
if let Ok(mut canonical) = std::fs::canonicalize(&ancestor) {
for component in missing_suffix.iter().rev() {
if component == std::ffi::OsStr::new(".") {
continue;
}
if component == std::ffi::OsStr::new("..") {
canonical.pop();
} else {
canonical.push(component);
}
}
return canonical;
}
}
absolute
}
pub fn install_shared_local_admission(coordinator: Arc<LocalAdmissionCoordinator>) {
let root = normalized_state_root_key(&car_home::root_or_relative());
scoped_admission_registry()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.entry(root)
.or_insert_with(|| Arc::downgrade(&coordinator));
}
pub fn shared_local_admission() -> Arc<LocalAdmissionCoordinator> {
let root = car_home::root_or_relative();
let policy = FileResourcePolicyRepository::new(root.clone())
.load()
.unwrap_or_else(|_| ResourcePolicy::everyday());
scoped_local_admission(root, policy, HardwareInfo::detect())
}
pub fn scoped_local_admission(
state_root: impl AsRef<Path>,
policy: ResourcePolicy,
hardware: HardwareInfo,
) -> Arc<LocalAdmissionCoordinator> {
let state_root = normalized_state_root_key(state_root.as_ref());
let mut registry = scoped_admission_registry()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
registry.retain(|_, coordinator| coordinator.strong_count() > 0);
if let Some(existing) = registry.get(&state_root).and_then(Weak::upgrade) {
existing.set_policy(policy);
return existing;
}
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
policy,
hardware,
Arc::new(SystemLiveMemoryProbe),
process_machine_admission_ledger(),
));
registry.insert(state_root, Arc::downgrade(&coordinator));
coordinator
}
pub fn local_admission_for_scope(
state_root: impl AsRef<Path>,
) -> Option<Arc<LocalAdmissionCoordinator>> {
let state_root = normalized_state_root_key(state_root.as_ref());
scoped_admission_registry()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&state_root)
.and_then(Weak::upgrade)
}
impl LocalAdmissionCoordinator {
fn built_in_model_identity(model_id: &str) -> Option<String> {
let lower = model_id.to_ascii_lowercase();
if lower.contains("kokoro-82m") {
if lower.contains("6bit") {
return Some("mlx/kokoro-82m:6bit".into());
}
if lower.contains("bf16") {
return Some("mlx/kokoro-82m:bf16".into());
}
}
None
}
fn resolve_model_identity(state: &AdmissionState, model_id: &str) -> String {
state
.model_aliases
.iter()
.find_map(|(canonical, aliases)| {
(canonical == model_id || aliases.contains(model_id)).then(|| canonical.clone())
})
.or_else(|| Self::built_in_model_identity(model_id))
.unwrap_or_else(|| model_id.to_string())
}
pub fn canonical_model_id(&self, model_id: &str) -> String {
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
Self::resolve_model_identity(&state, model_id)
}
pub fn register_model_aliases<I, S>(&self, canonical_model_id: &str, aliases: I)
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
let mut aliases = aliases
.into_iter()
.map(Into::into)
.collect::<HashSet<String>>();
aliases.insert(canonical_model_id.to_string());
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut canonical_model_id = Self::built_in_model_identity(canonical_model_id)
.unwrap_or_else(|| canonical_model_id.to_string());
let intersecting = state
.model_aliases
.iter()
.filter(|(existing, existing_aliases)| {
aliases.contains(*existing)
|| existing_aliases.iter().any(|alias| aliases.contains(alias))
})
.map(|(existing, _)| existing.clone())
.collect::<Vec<_>>();
if let Some(existing) = intersecting.first() {
canonical_model_id = existing.clone();
}
for existing in intersecting {
if let Some(existing_aliases) = state.model_aliases.remove(&existing) {
aliases.extend(existing_aliases);
}
aliases.insert(existing);
}
aliases.insert(canonical_model_id.clone());
let mut active = 0usize;
for alias in &aliases {
active = active.saturating_add(state.active_by_model.remove(alias).unwrap_or_default());
}
if active > 0 {
*state
.active_by_model
.entry(canonical_model_id.clone())
.or_default() += active;
}
let mut maintenance = false;
for alias in &aliases {
maintenance |= state.maintenance_models.remove(alias);
}
if maintenance {
state.maintenance_models.insert(canonical_model_id.clone());
}
let mut pending = HashSet::new();
for alias in &aliases {
pending.extend(
state
.pending_teardown_models
.remove(alias)
.unwrap_or_default(),
);
}
if !pending.is_empty() {
state
.pending_teardown_models
.entry(canonical_model_id.clone())
.or_default()
.extend(pending);
}
for resident in state.resident_models.values_mut() {
if aliases.contains(&resident.logical_model_id) {
resident.logical_model_id = canonical_model_id.clone();
}
}
state
.model_aliases
.insert(canonical_model_id.clone(), aliases.clone());
let mut machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
for ((owner_id, _), resident) in machine.resident_models.iter_mut() {
if *owner_id == self.owner_id && aliases.contains(&resident.logical_model_id) {
resident.logical_model_id = canonical_model_id.clone();
}
}
for ((owner_id, _), resident) in machine.pending_allocations.iter_mut() {
if *owner_id == self.owner_id && aliases.contains(&resident.logical_model_id) {
resident.logical_model_id = canonical_model_id.clone();
}
}
}
pub fn new(policy: ResourcePolicy, hardware: HardwareInfo) -> Self {
Self::with_probe(policy, hardware, Arc::new(SystemLiveMemoryProbe))
}
pub fn with_probe(
policy: ResourcePolicy,
hardware: HardwareInfo,
live_probe: Arc<dyn LiveMemoryProbe>,
) -> Self {
Self::with_probe_and_ledger(
policy,
hardware,
live_probe,
Arc::new(Mutex::new(MachineAdmissionLedger::default())),
)
}
fn with_probe_and_ledger(
policy: ResourcePolicy,
hardware: HardwareInfo,
live_probe: Arc<dyn LiveMemoryProbe>,
machine_ledger: Arc<Mutex<MachineAdmissionLedger>>,
) -> Self {
Self {
policy: std::sync::RwLock::new(policy),
hardware,
live_probe,
state: Mutex::new(AdmissionState::default()),
machine_ledger,
resident_activity_leases: Mutex::new(HashMap::new()),
owner_id: next_admission_owner_id(),
}
}
pub fn set_policy(&self, policy: ResourcePolicy) {
*self
.policy
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) = policy;
}
pub fn policy(&self) -> ResourcePolicy {
self.policy
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
pub fn mark_resident(&self, model_id: &str, weights_mb: u64) {
let placement = self.default_weight_placement();
self.mark_resident_with_placement(model_id, model_id, weights_mb, placement);
}
pub fn mark_resident_allocation(
&self,
logical_model_id: &str,
allocation_id: &str,
weights_mb: u64,
) {
let placement = self.default_weight_placement();
self.mark_resident_with_placement(logical_model_id, allocation_id, weights_mb, placement);
}
pub fn mark_evicted(&self, model_id: &str) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.resident_models.remove(model_id);
self.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.resident_models
.remove(&(self.owner_id, model_id.to_string()));
self.resident_activity_leases
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(model_id);
}
pub fn mark_teardown_pending(&self, model_id: &str) {
self.mark_teardown_pending_allocation(model_id, model_id);
}
pub fn mark_teardown_pending_allocation(&self, logical_model_id: &str, allocation_id: &str) {
self.mark_teardown_pending_allocation_with_charge(logical_model_id, allocation_id, 0);
}
pub fn mark_teardown_pending_allocation_with_charge(
&self,
logical_model_id: &str,
allocation_id: &str,
measured_bytes: u64,
) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, logical_model_id);
state
.pending_teardown_models
.entry(model_id.clone())
.or_default()
.insert(allocation_id.to_string());
let mut machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let key = (self.owner_id, allocation_id.to_string());
let machine_resident = machine.resident_models.remove(&key);
let scoped_resident = state.resident_models.get(allocation_id).cloned();
let resident = machine_resident.or(scoped_resident);
let measured_mb = measured_bytes.div_ceil(1024 * 1024);
let placement = self.default_weight_placement();
machine
.pending_allocations
.entry(key)
.and_modify(|pending| {
pending.weights_mb = pending.weights_mb.max(measured_mb);
pending.logical_model_id = model_id.clone();
})
.or_insert_with(|| {
resident.unwrap_or(ResidentAllocation {
weights_mb: measured_mb,
placement,
logical_model_id: model_id,
})
});
}
pub fn finish_teardown(&self, model_id: &str) {
self.finish_teardown_allocation(model_id, model_id);
}
pub fn finish_teardown_allocation(&self, logical_model_id: &str, allocation_id: &str) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
let exact_pending = state
.pending_teardown_models
.get_mut(&logical_model_id)
.is_some_and(|pending| pending.remove(allocation_id));
if exact_pending
&& state
.pending_teardown_models
.get(&logical_model_id)
.is_some_and(HashSet::is_empty)
{
state.pending_teardown_models.remove(&logical_model_id);
}
if !exact_pending {
return;
}
state.resident_models.remove(allocation_id);
let mut machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
machine
.resident_models
.remove(&(self.owner_id, allocation_id.to_string()));
machine
.pending_allocations
.remove(&(self.owner_id, allocation_id.to_string()));
self.resident_activity_leases
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(allocation_id);
}
pub fn teardown_pending(&self, model_id: &str) -> bool {
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, model_id);
state.pending_teardown_models.contains_key(&model_id)
}
pub fn is_resident(&self, model_id: &str) -> bool {
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, model_id);
state
.resident_models
.iter()
.any(|(allocation_id, resident)| {
allocation_id == &model_id || resident.logical_model_id == model_id
})
}
pub fn resident_model_mb(&self) -> u64 {
let machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
machine
.resident_models
.values()
.chain(machine.pending_allocations.values())
.map(|resident| resident.weights_mb)
.fold(0, u64::saturating_add)
}
pub fn active_request_count(&self, model_id: &str) -> usize {
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, model_id);
state
.active_by_model
.get(&model_id)
.copied()
.unwrap_or_default()
}
pub fn resident_allocation_ids(&self, model_id: &str) -> Vec<String> {
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, model_id);
let mut allocations = state
.resident_models
.iter()
.filter(|(allocation_id, resident)| {
allocation_id.as_str() == model_id || resident.logical_model_id == model_id
})
.map(|(allocation_id, _)| allocation_id.clone())
.collect::<Vec<_>>();
if let Some(pending) = state.pending_teardown_models.get(&model_id) {
allocations.extend(pending.iter().cloned());
}
allocations.sort();
allocations.dedup();
allocations
}
fn default_weight_placement(&self) -> WeightPlacement {
if matches!(self.hardware.gpu_backend, GpuBackend::Cuda)
&& self.hardware.gpu_memory_mb.is_some()
{
WeightPlacement::Accelerator
} else {
WeightPlacement::Host
}
}
fn mark_resident_with_placement(
&self,
logical_model_id: &str,
allocation_id: &str,
weights_mb: u64,
placement: WeightPlacement,
) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
let allocation = ResidentAllocation {
weights_mb,
placement,
logical_model_id: logical_model_id.clone(),
};
state
.resident_models
.entry(allocation_id.to_string())
.and_modify(|resident| {
resident.weights_mb = resident.weights_mb.max(weights_mb);
resident.placement = placement;
resident.logical_model_id = logical_model_id.clone();
})
.or_insert(ResidentAllocation {
weights_mb,
placement,
logical_model_id,
});
self.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.resident_models
.insert((self.owner_id, allocation_id.to_string()), allocation);
}
pub fn preflight(&self, model: &ModelSchema, context_tokens: usize) -> LocalLoadPreflight {
let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
let state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
self.preflight_locked(model, context_tokens, &state, &machine, live_available_mb)
}
pub fn reserve(
self: &Arc<Self>,
model: &ModelSchema,
context_tokens: usize,
) -> Result<LocalLoadReservation, LocalAdmissionError> {
let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, &model.id);
let estimate = estimate_model_memory(model, &self.hardware, context_tokens);
let mut preflight = self.preflight_estimate_locked(
&model_id,
estimate,
self.default_weight_placement(),
&state,
&machine,
live_available_mb,
);
if state.pending_teardown_models.contains_key(&model_id) {
preflight.verdict = LocalLoadVerdict::PendingTeardown;
return Err(LocalAdmissionError { preflight });
}
if state.maintenance_models.contains(&model_id) {
preflight.verdict = LocalLoadVerdict::ModelMaintenance;
return Err(LocalAdmissionError { preflight });
}
if !preflight.verdict.permits_static_fallback() {
return Err(LocalAdmissionError { preflight });
}
if preflight.verdict == LocalLoadVerdict::LiveMemoryUnknown {
tracing::warn!(
model = %model.id,
configured_ceiling_mb = preflight.configured_ceiling_mb,
estimated_incremental_mb = preflight.estimated_incremental_mb,
"live memory is unknown; proceeding under the static configured ceiling only"
);
}
state.active_host_reservations_mb = state
.active_host_reservations_mb
.saturating_add(preflight.estimated_incremental_mb);
state.active_accelerator_reservations_mb = state
.active_accelerator_reservations_mb
.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
let host = machine
.active_host_by_owner
.entry(self.owner_id)
.or_default();
*host = host.saturating_add(preflight.estimated_incremental_mb);
let accelerator = machine
.active_accelerator_by_owner
.entry(self.owner_id)
.or_default();
*accelerator =
accelerator.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
*state.active_by_model.entry(model_id.clone()).or_default() += 1;
state.next_request_id = state.next_request_id.wrapping_add(1);
let request_id = state.next_request_id;
Ok(LocalLoadReservation {
request_id,
model_id: model_id.clone(),
maintenance_model_id: model_id.clone(),
weights_mb: preflight.estimate.weights_mb,
reserved_incremental_mb: preflight.estimated_incremental_mb,
reserved_accelerator_mb: preflight.accelerator_incremental_mb.unwrap_or_default(),
cold_weights_reserved: !Self::has_resident_model(&state, &model_id),
placement: self.default_weight_placement(),
coordinator: Arc::clone(self),
charge: Arc::new(ReservationCharge::new(
model_id,
preflight.estimated_incremental_mb,
preflight.accelerator_incremental_mb.unwrap_or_default(),
Arc::clone(self),
)),
})
}
pub fn reserve_measured_host(
self: &Arc<Self>,
model_id: &str,
measured_weights_bytes: u64,
request_overhead_mb: u64,
) -> Result<LocalLoadReservation, LocalAdmissionError> {
self.reserve_measured(
model_id,
model_id,
measured_weights_bytes,
request_overhead_mb,
WeightPlacement::Host,
)
}
pub fn reserve_measured_host_allocation(
self: &Arc<Self>,
logical_model_id: &str,
allocation_id: &str,
measured_weights_bytes: u64,
request_overhead_mb: u64,
) -> Result<LocalLoadReservation, LocalAdmissionError> {
self.reserve_measured(
logical_model_id,
allocation_id,
measured_weights_bytes,
request_overhead_mb,
WeightPlacement::Host,
)
}
fn reserve_measured(
self: &Arc<Self>,
logical_model_id: &str,
allocation_id: &str,
measured_weights_bytes: u64,
request_overhead_mb: u64,
placement: WeightPlacement,
) -> Result<LocalLoadReservation, LocalAdmissionError> {
let weights_mb = measured_weights_bytes.div_ceil(1024 * 1024);
let estimate = ModelMemoryEstimate {
weights_mb,
runtime_overhead_mb: request_overhead_mb,
context_overhead_mb: 0,
transient_margin_mb: 0,
estimated_peak_mb: weights_mb.saturating_add(request_overhead_mb),
evidence: ModelResourceEvidence::FileSystemMeasured,
};
let live_available_mb = self.live_probe.available_memory_mb().ok().flatten();
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut machine = self
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let logical_model_id = Self::resolve_model_identity(&state, logical_model_id);
let mut preflight = self.preflight_estimate_locked(
allocation_id,
estimate,
placement,
&state,
&machine,
live_available_mb,
);
if state
.pending_teardown_models
.contains_key(&logical_model_id)
{
preflight.verdict = LocalLoadVerdict::PendingTeardown;
}
if state.maintenance_models.contains(&logical_model_id) {
preflight.verdict = LocalLoadVerdict::ModelMaintenance;
}
if !preflight.verdict.permits_static_fallback() {
return Err(LocalAdmissionError { preflight });
}
state.active_host_reservations_mb = state
.active_host_reservations_mb
.saturating_add(preflight.estimated_incremental_mb);
state.active_accelerator_reservations_mb = state
.active_accelerator_reservations_mb
.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
let host = machine
.active_host_by_owner
.entry(self.owner_id)
.or_default();
*host = host.saturating_add(preflight.estimated_incremental_mb);
let accelerator = machine
.active_accelerator_by_owner
.entry(self.owner_id)
.or_default();
*accelerator =
accelerator.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
*state
.active_by_model
.entry(logical_model_id.clone())
.or_default() += 1;
state.next_request_id = state.next_request_id.wrapping_add(1);
Ok(LocalLoadReservation {
request_id: state.next_request_id,
model_id: allocation_id.to_string(),
maintenance_model_id: logical_model_id.clone(),
weights_mb,
reserved_incremental_mb: preflight.estimated_incremental_mb,
reserved_accelerator_mb: preflight.accelerator_incremental_mb.unwrap_or_default(),
cold_weights_reserved: !state.resident_models.contains_key(allocation_id),
placement,
coordinator: Arc::clone(self),
charge: Arc::new(ReservationCharge::new(
logical_model_id,
preflight.estimated_incremental_mb,
preflight.accelerator_incremental_mb.unwrap_or_default(),
Arc::clone(self),
)),
})
}
pub fn begin_model_maintenance(
self: &Arc<Self>,
model_id: &str,
) -> Result<LocalModelMaintenanceGuard, ModelMaintenanceError> {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let model_id = Self::resolve_model_identity(&state, model_id);
if state
.active_by_model
.get(&model_id)
.copied()
.unwrap_or_default()
> 0
{
return Err(ModelMaintenanceError::ModelInUse(model_id));
}
if !state.maintenance_models.insert(model_id.clone()) {
return Err(ModelMaintenanceError::AlreadyInMaintenance(model_id));
}
Ok(LocalModelMaintenanceGuard {
model_id,
coordinator: Arc::clone(self),
})
}
fn preflight_locked(
&self,
model: &ModelSchema,
context_tokens: usize,
state: &AdmissionState,
machine: &MachineAdmissionLedger,
live_available_mb: Option<u64>,
) -> LocalLoadPreflight {
let estimate = estimate_model_memory(model, &self.hardware, context_tokens);
let model_id = Self::resolve_model_identity(state, &model.id);
self.preflight_estimate_locked(
&model_id,
estimate,
self.default_weight_placement(),
state,
machine,
live_available_mb,
)
}
fn preflight_estimate_locked(
&self,
model_id: &str,
estimate: ModelMemoryEstimate,
placement: WeightPlacement,
state: &AdmissionState,
machine: &MachineAdmissionLedger,
live_available_mb: Option<u64>,
) -> LocalLoadPreflight {
let policy = self.policy();
let budget = policy.effective_budget(self.hardware.total_ram_mb);
let host_resident_mb = machine
.resident_models
.values()
.chain(machine.pending_allocations.values())
.filter(|resident| resident.placement == WeightPlacement::Host)
.map(|resident| resident.weights_mb)
.fold(0, u64::saturating_add);
let accelerator_resident_weights_mb = machine
.resident_models
.values()
.chain(machine.pending_allocations.values())
.filter(|resident| resident.placement == WeightPlacement::Accelerator)
.map(|resident| resident.weights_mb)
.fold(0, u64::saturating_add);
let already_resident = Self::has_resident_model(state, model_id);
let request_overhead_mb = estimate
.context_overhead_mb
.saturating_add(estimate.runtime_overhead_mb)
.saturating_add(estimate.transient_margin_mb);
let cold_weights_mb = if already_resident {
0
} else {
estimate.weights_mb
};
let resident_model_mb = host_resident_mb;
let host_cold_weights_mb = if placement == WeightPlacement::Host {
cold_weights_mb
} else {
0
};
let estimated_incremental_mb = request_overhead_mb.saturating_add(host_cold_weights_mb);
let accelerator_total_mb = (placement == WeightPlacement::Accelerator)
.then_some(self.hardware.gpu_memory_mb)
.flatten();
let accelerator_resident_mb =
(placement == WeightPlacement::Accelerator).then_some(accelerator_resident_weights_mb);
let accelerator_incremental_mb =
(placement == WeightPlacement::Accelerator).then_some(cold_weights_mb);
let projected_static_mb = resident_model_mb
.saturating_add(machine.active_host_by_owner.values().copied().sum::<u64>())
.saturating_add(estimated_incremental_mb);
let projected_accelerator_mb = accelerator_resident_weights_mb
.saturating_add(
machine
.active_accelerator_by_owner
.values()
.copied()
.sum::<u64>(),
)
.saturating_add(cold_weights_mb);
let verdict = if budget.effective_new_load_ceiling_mb == 0 && !already_resident {
LocalLoadVerdict::DisabledByPolicy
} else if (budget.configured_model_ceiling_mb != 0
&& projected_static_mb > budget.configured_model_ceiling_mb)
|| accelerator_total_mb.is_some_and(|vram_mb| projected_accelerator_mb > vram_mb)
{
LocalLoadVerdict::ExceedsConfiguredCeiling
} else if let Some(available_mb) = live_available_mb {
let unreserved_available_mb = available_mb
.saturating_sub(machine.active_host_by_owner.values().copied().sum::<u64>());
if unreserved_available_mb
< estimated_incremental_mb.saturating_add(budget.emergency_reserve_mb)
{
LocalLoadVerdict::InsufficientLiveMemory
} else {
LocalLoadVerdict::Allowed
}
} else {
LocalLoadVerdict::LiveMemoryUnknown
};
LocalLoadPreflight {
model_id: model_id.to_string(),
estimate,
configured_ceiling_mb: budget.configured_model_ceiling_mb,
resident_model_mb,
active_reservations_mb: machine.active_host_by_owner.values().copied().sum(),
estimated_incremental_mb,
accelerator_total_mb,
accelerator_resident_mb,
accelerator_incremental_mb,
live_available_mb,
emergency_reserve_mb: budget.emergency_reserve_mb,
verdict,
}
}
fn has_resident_model(state: &AdmissionState, model_id: &str) -> bool {
state
.resident_models
.iter()
.any(|(allocation_id, resident)| {
allocation_id == model_id || resident.logical_model_id == model_id
})
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LocalAdmissionError {
pub preflight: LocalLoadPreflight,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum ModelMaintenanceError {
#[error("local model '{0}' is in use")]
ModelInUse(String),
#[error("local model '{0}' already has maintenance in progress")]
AlreadyInMaintenance(String),
#[error("failed to release local model residency: {0}")]
ReleaseFailed(String),
#[error("local worker did not acknowledge release of model '{0}'")]
WorkerReleaseUnacknowledged(String),
#[error("supervised local process did not acknowledge release of model '{0}'")]
ProcessReleaseUnacknowledged(String),
#[error("an in-process cache still has active work for local model '{0}'")]
CacheReleaseBlocked(String),
#[error("local model '{model_id}' still has resident allocations: {allocation_ids:?}")]
ResidualResidency {
model_id: String,
allocation_ids: Vec<String>,
},
}
pub struct LocalModelMaintenanceGuard {
model_id: String,
coordinator: Arc<LocalAdmissionCoordinator>,
}
impl Drop for LocalModelMaintenanceGuard {
fn drop(&mut self) {
self.coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.maintenance_models
.remove(&self.model_id);
}
}
impl std::fmt::Display for LocalAdmissionError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"local model '{}' blocked by {:?}: needs {} MB incremental, {} MB live available",
self.preflight.model_id,
self.preflight.verdict,
self.preflight.estimated_incremental_mb,
self.preflight
.live_available_mb
.map(|value| value.to_string())
.unwrap_or_else(|| "unknown".into())
)
}
}
impl std::error::Error for LocalAdmissionError {}
pub struct LocalLoadReservation {
request_id: u64,
model_id: String,
maintenance_model_id: String,
weights_mb: u64,
reserved_incremental_mb: u64,
reserved_accelerator_mb: u64,
cold_weights_reserved: bool,
placement: WeightPlacement,
coordinator: Arc<LocalAdmissionCoordinator>,
charge: Arc<ReservationCharge>,
}
#[derive(Clone)]
pub struct DetachedLocalLease {
_charge: Arc<ReservationCharge>,
}
struct ReservationCharge {
maintenance_model_id: String,
host_mb: AtomicU64,
accelerator_mb: AtomicU64,
cold_weights_transferred: AtomicBool,
coordinator: Arc<LocalAdmissionCoordinator>,
activity_lease: Mutex<Option<Arc<crate::model_management::ModelLease>>>,
}
impl ReservationCharge {
fn new(
maintenance_model_id: String,
host_mb: u64,
accelerator_mb: u64,
coordinator: Arc<LocalAdmissionCoordinator>,
) -> Self {
Self {
maintenance_model_id,
host_mb: AtomicU64::new(host_mb),
accelerator_mb: AtomicU64::new(accelerator_mb),
cold_weights_transferred: AtomicBool::new(false),
coordinator,
activity_lease: Mutex::new(None),
}
}
fn update(&self, host_mb: u64, accelerator_mb: u64) {
self.host_mb.store(host_mb, Ordering::Release);
self.accelerator_mb.store(accelerator_mb, Ordering::Release);
}
}
impl Drop for ReservationCharge {
fn drop(&mut self) {
let host_mb = self.host_mb.load(Ordering::Acquire);
let accelerator_mb = self.accelerator_mb.load(Ordering::Acquire);
let mut state = self
.coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut machine = self
.coordinator
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.active_host_reservations_mb =
state.active_host_reservations_mb.saturating_sub(host_mb);
state.active_accelerator_reservations_mb = state
.active_accelerator_reservations_mb
.saturating_sub(accelerator_mb);
if let Some(host) = machine
.active_host_by_owner
.get_mut(&self.coordinator.owner_id)
{
*host = host.saturating_sub(host_mb);
}
if let Some(accelerator) = machine
.active_accelerator_by_owner
.get_mut(&self.coordinator.owner_id)
{
*accelerator = accelerator.saturating_sub(accelerator_mb);
}
let maintenance_model_id =
LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
if let Some(active) = state.active_by_model.get_mut(&maintenance_model_id) {
*active = active.saturating_sub(1);
if *active == 0 {
state.active_by_model.remove(&maintenance_model_id);
}
}
}
}
impl std::fmt::Debug for LocalLoadReservation {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("LocalLoadReservation")
.field("request_id", &self.request_id)
.field("model_id", &self.model_id)
.field("reserved_incremental_mb", &self.reserved_incremental_mb)
.field("reserved_accelerator_mb", &self.reserved_accelerator_mb)
.finish_non_exhaustive()
}
}
impl LocalLoadReservation {
pub fn request_id(&self) -> u64 {
self.request_id
}
pub fn model_id(&self) -> &str {
&self.model_id
}
pub fn authorizes_model(&self, model_id: &str) -> bool {
self.model_id == model_id || self.maintenance_model_id == model_id
}
pub(crate) fn bind_allocation_id(&mut self, allocation_id: &str) {
self.model_id = allocation_id.to_string();
}
pub fn reserved_incremental_mb(&self) -> u64 {
self.reserved_incremental_mb
}
pub fn reconciled_weights_bytes(&self) -> u64 {
self.weights_mb.saturating_mul(1024 * 1024)
}
pub fn detached_lease(&self) -> DetachedLocalLease {
DetachedLocalLease {
_charge: self.charge.clone(),
}
}
pub(crate) fn attach_activity_lease(&mut self, lease: crate::model_management::ModelLease) {
*self
.charge
.activity_lease
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::new(lease));
}
pub(crate) fn transfer_cold_weights_to_pending_allocation(
&self,
allocation_id: &str,
measured_weights_bytes: u64,
) {
let mut state = self
.coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let logical_model_id =
LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
state
.pending_teardown_models
.entry(logical_model_id.clone())
.or_default()
.insert(allocation_id.to_string());
let mut machine = self
.coordinator
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let key = (self.coordinator.owner_id, allocation_id.to_string());
let machine_resident = machine.resident_models.remove(&key);
let scoped_resident = state.resident_models.get(allocation_id).cloned();
let resident = machine_resident.or(scoped_resident);
let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
machine
.pending_allocations
.entry(key)
.and_modify(|pending| {
pending.weights_mb = pending.weights_mb.max(measured_mb);
pending.logical_model_id = logical_model_id.clone();
})
.or_insert_with(|| {
resident.unwrap_or(ResidentAllocation {
weights_mb: measured_mb,
placement: self.placement,
logical_model_id,
})
});
if self.cold_weights_reserved
&& self
.charge
.cold_weights_transferred
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
match self.placement {
WeightPlacement::Host => {
state.active_host_reservations_mb = state
.active_host_reservations_mb
.saturating_sub(self.weights_mb);
let active = machine
.active_host_by_owner
.entry(self.coordinator.owner_id)
.or_default();
*active = active.saturating_sub(self.weights_mb);
let _ = self.charge.host_mb.fetch_update(
Ordering::AcqRel,
Ordering::Acquire,
|host_mb| Some(host_mb.saturating_sub(self.weights_mb)),
);
}
WeightPlacement::Accelerator => {
state.active_accelerator_reservations_mb = state
.active_accelerator_reservations_mb
.saturating_sub(self.weights_mb);
let active = machine
.active_accelerator_by_owner
.entry(self.coordinator.owner_id)
.or_default();
*active = active.saturating_sub(self.weights_mb);
let _ = self.charge.accelerator_mb.fetch_update(
Ordering::AcqRel,
Ordering::Acquire,
|accelerator_mb| Some(accelerator_mb.saturating_sub(self.weights_mb)),
);
}
}
}
}
fn sync_shared_charge(&self) {
self.charge
.update(self.reserved_incremental_mb, self.reserved_accelerator_mb);
}
pub fn reconcile_measured_weights(
&mut self,
measured_weights_bytes: u64,
) -> Result<LocalLoadPreflight, LocalAdmissionError> {
if !self.cold_weights_reserved {
let live_available_mb = self
.coordinator
.live_probe
.available_memory_mb()
.ok()
.flatten();
let state = self
.coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let machine = self
.coordinator
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if LocalAdmissionCoordinator::has_resident_model(&state, &self.model_id) {
let estimate = self.measured_estimate(self.weights_mb);
return Ok(self.coordinator.preflight_estimate_locked(
&self.model_id,
estimate,
self.placement,
&state,
&machine,
live_available_mb,
));
}
drop(state);
self.cold_weights_reserved = true;
}
let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
let live_available_mb = self
.coordinator
.live_probe
.available_memory_mb()
.ok()
.flatten();
let mut state = self
.coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut machine = self
.coordinator
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let published_by_peer =
LocalAdmissionCoordinator::has_resident_model(&state, &self.model_id);
let mut without_current = state.clone();
without_current.active_host_reservations_mb = without_current
.active_host_reservations_mb
.saturating_sub(self.reserved_incremental_mb);
without_current.active_accelerator_reservations_mb = without_current
.active_accelerator_reservations_mb
.saturating_sub(self.reserved_accelerator_mb);
let mut without_machine = machine.clone();
let host = without_machine
.active_host_by_owner
.entry(self.coordinator.owner_id)
.or_default();
*host = host.saturating_sub(self.reserved_incremental_mb);
let accelerator = without_machine
.active_accelerator_by_owner
.entry(self.coordinator.owner_id)
.or_default();
*accelerator = accelerator.saturating_sub(self.reserved_accelerator_mb);
let estimate = self.measured_estimate(measured_mb);
let preflight = self.coordinator.preflight_estimate_locked(
&self.model_id,
estimate,
self.placement,
&without_current,
&without_machine,
live_available_mb,
);
if !preflight.verdict.permits_static_fallback() {
return Err(LocalAdmissionError { preflight });
}
state.active_host_reservations_mb = without_current
.active_host_reservations_mb
.saturating_add(preflight.estimated_incremental_mb);
state.active_accelerator_reservations_mb = without_current
.active_accelerator_reservations_mb
.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default());
machine.active_host_by_owner.insert(
self.coordinator.owner_id,
without_machine
.active_host_by_owner
.get(&self.coordinator.owner_id)
.copied()
.unwrap_or_default()
.saturating_add(preflight.estimated_incremental_mb),
);
machine.active_accelerator_by_owner.insert(
self.coordinator.owner_id,
without_machine
.active_accelerator_by_owner
.get(&self.coordinator.owner_id)
.copied()
.unwrap_or_default()
.saturating_add(preflight.accelerator_incremental_mb.unwrap_or_default()),
);
self.weights_mb = measured_mb;
self.reserved_incremental_mb = preflight.estimated_incremental_mb;
self.reserved_accelerator_mb = preflight.accelerator_incremental_mb.unwrap_or_default();
if published_by_peer {
self.cold_weights_reserved = false;
}
self.sync_shared_charge();
Ok(preflight)
}
fn measured_estimate(&self, weights_mb: u64) -> ModelMemoryEstimate {
let host_cold_mb = if self.cold_weights_reserved && self.placement == WeightPlacement::Host
{
self.weights_mb
} else {
0
};
let request_overhead_mb = self.reserved_incremental_mb.saturating_sub(host_cold_mb);
ModelMemoryEstimate {
weights_mb,
runtime_overhead_mb: request_overhead_mb,
context_overhead_mb: 0,
transient_margin_mb: 0,
estimated_peak_mb: weights_mb.saturating_add(request_overhead_mb),
evidence: ModelResourceEvidence::FileSystemMeasured,
}
}
pub fn publish_resident_weights(&mut self, measured_weights_bytes: u64) {
let allocation_id = self.model_id.clone();
self.publish_resident_weights_as(&allocation_id, measured_weights_bytes);
}
pub fn publish_resident_weights_as(
&mut self,
allocation_id: &str,
measured_weights_bytes: u64,
) {
let measured_mb = measured_weights_bytes.div_ceil(1024 * 1024);
let mut state = self
.coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut machine = self
.coordinator
.machine_ledger
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let key = (self.coordinator.owner_id, allocation_id.to_string());
let pending = machine.pending_allocations.remove(&key);
let logical_model_id =
LocalAdmissionCoordinator::resolve_model_identity(&state, &self.maintenance_model_id);
if let Some(allocations) = state.pending_teardown_models.get_mut(&logical_model_id) {
allocations.remove(allocation_id);
if allocations.is_empty() {
state.pending_teardown_models.remove(&logical_model_id);
}
}
if !self.cold_weights_reserved && pending.is_none() {
return;
}
if self.cold_weights_reserved {
let cold_weights_transferred =
self.charge.cold_weights_transferred.load(Ordering::Acquire);
match self.placement {
WeightPlacement::Host => {
if !cold_weights_transferred {
state.active_host_reservations_mb = state
.active_host_reservations_mb
.saturating_sub(self.weights_mb);
let active = machine
.active_host_by_owner
.entry(self.coordinator.owner_id)
.or_default();
*active = active.saturating_sub(self.weights_mb);
}
self.reserved_incremental_mb =
self.reserved_incremental_mb.saturating_sub(self.weights_mb);
}
WeightPlacement::Accelerator => {
if !cold_weights_transferred {
state.active_accelerator_reservations_mb = state
.active_accelerator_reservations_mb
.saturating_sub(self.weights_mb);
let active = machine
.active_accelerator_by_owner
.entry(self.coordinator.owner_id)
.or_default();
*active = active.saturating_sub(self.weights_mb);
}
self.reserved_accelerator_mb =
self.reserved_accelerator_mb.saturating_sub(self.weights_mb);
}
}
}
let resident = ResidentAllocation {
weights_mb: pending
.as_ref()
.map(|allocation| allocation.weights_mb)
.unwrap_or_default()
.max(measured_mb),
placement: pending
.as_ref()
.map(|allocation| allocation.placement)
.unwrap_or(self.placement),
logical_model_id,
};
state
.resident_models
.insert(allocation_id.to_string(), resident.clone());
machine.resident_models.insert(key, resident.clone());
if let Some(lease) = self
.charge
.activity_lease
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
{
self.coordinator
.resident_activity_leases
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(allocation_id.to_string(), lease);
}
self.weights_mb = resident.weights_mb;
self.cold_weights_reserved = false;
self.charge
.cold_weights_transferred
.store(false, Ordering::Release);
self.sync_shared_charge();
}
pub fn commit_resident_weights(&mut self) {
self.publish_resident_weights(self.weights_mb * 1024 * 1024);
}
}
pub fn estimate_model_memory(
model: &ModelSchema,
hardware: &HardwareInfo,
context_tokens: usize,
) -> ModelMemoryEstimate {
estimate_model_memory_with_measured_weights(model, hardware, context_tokens, None)
}
pub fn estimate_model_memory_with_measured_weights(
model: &ModelSchema,
hardware: &HardwareInfo,
context_tokens: usize,
measured_weights_mb: Option<u64>,
) -> ModelMemoryEstimate {
let declared = model.cost.ram_mb.or(model.cost.size_mb);
let (weights_mb, evidence) = if let Some(measured) = measured_weights_mb {
(measured, ModelResourceEvidence::FileSystemMeasured)
} else if let Some(declared) = declared {
(
declared.max(model.cost.size_mb.unwrap_or(0)),
ModelResourceEvidence::CatalogExact,
)
} else {
(
heuristic_weights_mb(model),
ModelResourceEvidence::Heuristic,
)
};
let context_overhead_mb = kv_cache_mb(model, context_tokens);
let runtime_overhead_mb = backend_runtime_overhead_mb(hardware);
let transient_margin_mb = TRANSIENT_ALLOCATION_MARGIN_MB;
let estimated_peak_mb = weights_mb
.saturating_add(context_overhead_mb)
.saturating_add(runtime_overhead_mb)
.saturating_add(transient_margin_mb);
ModelMemoryEstimate {
weights_mb,
runtime_overhead_mb,
context_overhead_mb,
transient_margin_mb,
estimated_peak_mb,
evidence,
}
}
fn backend_runtime_overhead_mb(hardware: &HardwareInfo) -> u64 {
match hardware.gpu_backend {
GpuBackend::Metal => METAL_RUNTIME_OVERHEAD_MB,
GpuBackend::Cuda => CUDA_RUNTIME_OVERHEAD_MB,
_ => CPU_RUNTIME_OVERHEAD_MB,
}
}
fn kv_cache_mb(model: &ModelSchema, context_tokens: usize) -> u64 {
let per_1k = (model_parameter_billions_active(model) as f64 * 0.12).max(0.05);
((context_tokens as f64 / 1_000.0) * per_1k).ceil() as u64
}
fn heuristic_weights_mb(model: &ModelSchema) -> u64 {
let billions = model_parameter_billions_total(model);
(billions as f64 * 600.0).ceil() as u64
}
pub(crate) fn model_parameter_billions_active(model: &ModelSchema) -> f32 {
model
.param_count
.split_once('(')
.and_then(|(_, rest)| rest.split_once("active"))
.and_then(|(number, _)| parse_parameter_billions(number))
.unwrap_or_else(|| model_parameter_billions_total(model))
}
pub(crate) fn model_parameter_billions_total(model: &ModelSchema) -> f32 {
parse_parameter_billions(&model.param_count).unwrap_or_else(|| {
let size_mb = model.size_mb();
if size_mb > 0 {
(size_mb as f32 / 600.0).max(0.1)
} else {
0.0
}
})
}
pub(crate) fn parse_parameter_billions(value: &str) -> Option<f32> {
let value = value.trim();
let number: String = value
.chars()
.take_while(|character| character.is_ascii_digit() || *character == '.')
.collect();
let parsed: f32 = number.parse().ok()?;
if value[number.len()..]
.trim_start()
.to_ascii_lowercase()
.starts_with('m')
{
Some(parsed / 1_000.0)
} else {
Some(parsed)
}
}
pub fn evaluate_resources(hardware: &HardwareInfo, policy: &ResourcePolicy) -> ResourceEvaluation {
let accelerator_memory = match (hardware.gpu_backend.clone(), hardware.gpu_memory_mb) {
(GpuBackend::Cuda, Some(total_mb)) => Some(AcceleratorResourceBudget {
total_mb,
budget_mb: total_mb,
}),
_ => None,
};
ResourceEvaluation {
host_memory: policy.effective_budget(hardware.total_ram_mb),
accelerator_memory,
}
}
fn percent_of(total_mb: u64, percent: u64) -> u64 {
let value = (total_mb as u128).saturating_mul(percent as u128) / 100;
value.min(u64::MAX as u128) as u64
}
fn minimum_emergency_reserve(total_memory_mb: u64) -> u64 {
MINIMUM_EMERGENCY_RESERVE_MB.max(percent_of(total_memory_mb, EMERGENCY_RESERVE_PERCENT))
}
#[derive(Debug, Error)]
pub enum ResourcePolicyError {
#[error("resource policy I/O failed: {0}")]
Io(#[from] io::Error),
#[error("resource policy serialization failed: {0}")]
Serialization(#[from] serde_json::Error),
#[error("Custom model RAM must be finite, nonnegative, and a 0.5 GB increment; got {0}")]
InvalidCustomGigabytes(f64),
#[error("invalid resource policy: {reason}")]
InvalidPolicy { reason: String },
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ResourcePolicyLoadSource {
Loaded,
MissingDefault,
CorruptDefault,
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct ResourcePolicyLoadEvidence {
pub policy: ResourcePolicy,
pub source: ResourcePolicyLoadSource,
pub warning: Option<String>,
}
pub trait ResourcePolicyRepository: Send + Sync {
fn load(&self) -> Result<ResourcePolicy, ResourcePolicyError>;
fn save(&self, policy: &ResourcePolicy) -> Result<(), ResourcePolicyError>;
}
#[derive(Clone, Debug)]
pub struct FileResourcePolicyRepository {
root: PathBuf,
}
impl Default for FileResourcePolicyRepository {
fn default() -> Self {
Self::new(car_home::root_or_relative())
}
}
impl FileResourcePolicyRepository {
pub fn new(root: PathBuf) -> Self {
Self { root }
}
pub fn path(&self) -> PathBuf {
self.root.join(RESOURCE_POLICY_FILE)
}
pub fn load_with_evidence(&self) -> Result<ResourcePolicyLoadEvidence, ResourcePolicyError> {
let path = self.path();
let file = match open_resource_policy(&path) {
Ok(file) => file,
Err(error) if error.kind() == io::ErrorKind::NotFound => {
return Ok(ResourcePolicyLoadEvidence {
policy: ResourcePolicy::everyday(),
source: ResourcePolicyLoadSource::MissingDefault,
warning: None,
});
}
Err(error) => {
if std::fs::symlink_metadata(&path)
.is_ok_and(|metadata| !metadata.is_file() || metadata.file_type().is_symlink())
{
return Ok(corrupt_default(
"The saved resource policy could not be loaded because it is not a regular file.",
));
}
return Err(error.into());
}
};
let metadata = file.metadata()?;
if !metadata.is_file() {
return Ok(corrupt_default(
"The saved resource policy could not be loaded because it is not a regular file.",
));
}
if metadata.len() > MAX_POLICY_BYTES {
return Ok(corrupt_default(
"The saved resource policy could not be loaded because it exceeds the size limit.",
));
}
let mut raw = Vec::new();
file.take(MAX_POLICY_BYTES + 1).read_to_end(&mut raw)?;
if raw.len() as u64 > MAX_POLICY_BYTES {
return Ok(corrupt_default(
"The saved resource policy could not be loaded because it exceeds the size limit.",
));
}
let policy = match serde_json::from_slice::<ResourcePolicy>(&raw) {
Ok(policy) => policy,
Err(error) => {
return Ok(corrupt_default(format!(
"The saved resource policy could not be loaded: {error}"
)));
}
};
if let Err(error) = policy.validate() {
return Ok(corrupt_default(format!(
"The saved resource policy could not be loaded: {error}"
)));
}
Ok(ResourcePolicyLoadEvidence {
policy,
source: ResourcePolicyLoadSource::Loaded,
warning: None,
})
}
}
fn open_resource_policy(path: &Path) -> io::Result<std::fs::File> {
let mut options = std::fs::OpenOptions::new();
options.read(true);
#[cfg(unix)]
{
use std::os::unix::fs::OpenOptionsExt;
options.custom_flags(libc::O_NOFOLLOW);
}
options.open(path)
}
fn corrupt_default(warning: impl Into<String>) -> ResourcePolicyLoadEvidence {
ResourcePolicyLoadEvidence {
policy: ResourcePolicy::everyday(),
source: ResourcePolicyLoadSource::CorruptDefault,
warning: Some(warning.into()),
}
}
impl ResourcePolicyRepository for FileResourcePolicyRepository {
fn load(&self) -> Result<ResourcePolicy, ResourcePolicyError> {
Ok(self.load_with_evidence()?.policy)
}
fn save(&self, policy: &ResourcePolicy) -> Result<(), ResourcePolicyError> {
policy.validate()?;
let _guard = mutation_lock()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
ensure_private_directory(&self.root)?;
let temp_path = unique_temp_path(&self.root);
let result = (|| {
let body = serde_json::to_vec_pretty(policy)?;
let mut file = open_private_temp(&temp_path)?;
file.write_all(&body)?;
file.sync_all()?;
atomic_replace(&temp_path, &self.path())?;
car_secrets::harden_owner_only(&self.path());
sync_directory(&self.root)?;
Ok(())
})();
if result.is_err() {
let _ = std::fs::remove_file(&temp_path);
}
result
}
}
fn mutation_lock() -> &'static Mutex<()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
}
fn unique_temp_path(root: &Path) -> PathBuf {
static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(0);
let sequence = NEXT_TEMP_ID.fetch_add(1, Ordering::Relaxed);
let epoch_nanos = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or(0);
root.join(format!(
".{RESOURCE_POLICY_FILE}.{}.{}.{}.tmp",
std::process::id(),
epoch_nanos,
sequence
))
}
fn ensure_private_directory(path: &Path) -> io::Result<()> {
std::fs::create_dir_all(path)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
}
car_secrets::harden_owner_only(path);
Ok(())
}
#[cfg(unix)]
fn open_private_temp(path: &Path) -> io::Result<std::fs::File> {
use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
let file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.mode(0o600)
.open(path)?;
file.set_permissions(std::fs::Permissions::from_mode(0o600))?;
Ok(file)
}
#[cfg(not(unix))]
fn open_private_temp(path: &Path) -> io::Result<std::fs::File> {
let file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(path)?;
car_secrets::harden_owner_only(path);
Ok(file)
}
#[cfg(not(windows))]
fn atomic_replace(source: &Path, destination: &Path) -> io::Result<()> {
std::fs::rename(source, destination)
}
#[cfg(windows)]
fn atomic_replace(source: &Path, destination: &Path) -> io::Result<()> {
use std::os::windows::ffi::OsStrExt;
const MOVEFILE_REPLACE_EXISTING: u32 = 0x1;
const MOVEFILE_WRITE_THROUGH: u32 = 0x8;
#[link(name = "kernel32")]
unsafe extern "system" {
fn MoveFileExW(existing: *const u16, new: *const u16, flags: u32) -> i32;
}
let source = source
.as_os_str()
.encode_wide()
.chain(Some(0))
.collect::<Vec<_>>();
let destination = destination
.as_os_str()
.encode_wide()
.chain(Some(0))
.collect::<Vec<_>>();
let replaced = unsafe {
MoveFileExW(
source.as_ptr(),
destination.as_ptr(),
MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
)
};
if replaced == 0 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
}
#[cfg(unix)]
fn sync_directory(path: &Path) -> io::Result<()> {
std::fs::File::open(path)?.sync_all()
}
#[cfg(not(unix))]
fn sync_directory(_path: &Path) -> io::Result<()> {
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn hardware(
total_ram_mb: u64,
gpu_backend: crate::hardware::GpuBackend,
vram_mb: Option<u64>,
) -> crate::hardware::HardwareInfo {
crate::hardware::HardwareInfo {
os: "test".into(),
arch: "test".into(),
cpu_cores: 8,
total_ram_mb,
gpu_backend,
gpu_memory_mb: vram_mb,
gpu_devices: Vec::new(),
recommended_model: "fixture".into(),
recommended_context: 4_096,
max_model_mb: total_ram_mb,
}
}
#[test]
fn profiles_compute_exact_32_gb_budgets() {
let total = 32 * 1024;
assert_eq!(
ResourcePolicy::everyday()
.effective_budget(total)
.configured_model_ceiling_mb,
13_107
);
assert_eq!(
ResourcePolicy::local_focused()
.effective_budget(total)
.configured_model_ceiling_mb,
26_214
);
assert_eq!(
ResourcePolicy::custom_gb(12.5)
.unwrap()
.effective_budget(total)
.configured_model_ceiling_mb,
12_800
);
assert_eq!(
ResourcePolicy::everyday().recommendation_target_mb(total),
6_553
);
}
#[test]
fn model_memory_estimate_keeps_transient_margin_distinct_and_totals_exactly() {
let catalog = crate::registry::builtin_catalog();
let model = catalog
.iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let estimate = estimate_model_memory(
model,
&hardware(32 * 1024, GpuBackend::Metal, None),
RECOMMENDATION_CONTEXT_TOKENS,
);
assert_eq!(estimate.evidence, ModelResourceEvidence::CatalogExact);
assert_eq!(estimate.weights_mb, 2_400);
assert_eq!(estimate.runtime_overhead_mb, 512);
assert_eq!(estimate.transient_margin_mb, 1_024);
assert_eq!(
estimate.estimated_peak_mb,
estimate.weights_mb
+ estimate.context_overhead_mb
+ estimate.runtime_overhead_mb
+ estimate.transient_margin_mb
);
let measured = estimate_model_memory_with_measured_weights(
model,
&hardware(32 * 1024, GpuBackend::Metal, None),
RECOMMENDATION_CONTEXT_TOKENS,
Some(2_321),
);
assert_eq!(measured.weights_mb, 2_321);
assert_eq!(measured.evidence, ModelResourceEvidence::FileSystemMeasured);
}
#[test]
fn custom_zero_disables_loads_and_overlarge_values_clamp_below_emergency_reserve() {
let total = 32 * 1024;
assert_eq!(
ResourcePolicy::custom_gb(0.0)
.unwrap()
.effective_budget(total)
.effective_new_load_ceiling_mb,
0
);
let result = ResourcePolicy::custom_gb(99.0)
.unwrap()
.effective_budget(total);
assert_eq!(result.emergency_reserve_mb, 3_276);
assert_eq!(result.configured_model_ceiling_mb, total - 3_276);
assert!(result.normalization_notice.is_some());
}
#[test]
fn custom_gigabytes_reject_non_half_steps_negative_and_non_finite_values() {
for invalid in [10.3, -0.5, f64::INFINITY, f64::NEG_INFINITY, f64::NAN] {
assert!(
ResourcePolicy::custom_gb(invalid).is_err(),
"accepted {invalid:?}"
);
}
assert_eq!(
ResourcePolicy::custom_gb(10.5).unwrap().custom_max_model_mb,
Some(10_752)
);
}
#[test]
fn repository_round_trips_exact_half_gb_and_uses_private_atomic_files() {
let dir = tempfile::tempdir().unwrap();
let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
repository
.save(&ResourcePolicy::custom_gb(10.5).unwrap())
.unwrap();
assert_eq!(
repository.path(),
dir.path().join("model-resource-policy.json")
);
assert_eq!(repository.load().unwrap().custom_max_model_mb, Some(10_752));
assert_private_mode(&repository.path(), 0o600);
assert_private_mode(dir.path(), 0o700);
assert_no_temp_files(dir.path());
}
#[test]
fn repository_missing_or_corrupt_file_falls_back_without_deleting_source() {
let dir = tempfile::tempdir().unwrap();
let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
assert_eq!(repository.load().unwrap(), ResourcePolicy::everyday());
let corrupt = br#"{"profile":"custom","custom_max_model_mb":"broken"}"#;
std::fs::write(repository.path(), corrupt).unwrap();
assert_eq!(repository.load().unwrap(), ResourcePolicy::everyday());
assert_eq!(std::fs::read(repository.path()).unwrap(), corrupt);
}
#[test]
fn repository_rejects_invalid_policy_shapes_before_writing() {
let dir = tempfile::tempdir().unwrap();
let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
let invalid = [
ResourcePolicy {
profile: ResourceProfile::Custom,
custom_max_model_mb: None,
},
ResourcePolicy {
profile: ResourceProfile::Custom,
custom_max_model_mb: Some(1),
},
ResourcePolicy {
profile: ResourceProfile::Custom,
custom_max_model_mb: Some(513),
},
ResourcePolicy {
profile: ResourceProfile::Everyday,
custom_max_model_mb: Some(512),
},
ResourcePolicy {
profile: ResourceProfile::LocalFocused,
custom_max_model_mb: Some(512),
},
];
for policy in invalid {
let error = repository.save(&policy).unwrap_err();
assert!(matches!(error, ResourcePolicyError::InvalidPolicy { .. }));
assert!(!repository.path().exists());
}
}
#[test]
fn load_evidence_distinguishes_loaded_missing_and_corrupt_defaults() {
let dir = tempfile::tempdir().unwrap();
let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
let missing = repository.load_with_evidence().unwrap();
assert_eq!(missing.policy, ResourcePolicy::everyday());
assert_eq!(missing.source, ResourcePolicyLoadSource::MissingDefault);
assert!(missing.warning.is_none());
repository.save(&ResourcePolicy::local_focused()).unwrap();
let loaded = repository.load_with_evidence().unwrap();
assert_eq!(loaded.policy, ResourcePolicy::local_focused());
assert_eq!(loaded.source, ResourcePolicyLoadSource::Loaded);
assert!(loaded.warning.is_none());
let corrupt = br#"{"profile":"custom","custom_max_model_mb":"broken"}"#;
std::fs::write(repository.path(), corrupt).unwrap();
let recovered = repository.load_with_evidence().unwrap();
assert_eq!(recovered.policy, ResourcePolicy::everyday());
assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
assert!(recovered
.warning
.as_deref()
.is_some_and(|warning| { warning.contains("could not be loaded") }));
assert_eq!(std::fs::read(repository.path()).unwrap(), corrupt);
}
#[test]
fn corrupt_evidence_covers_unknown_fields_invalid_shapes_and_oversized_files() {
let dir = tempfile::tempdir().unwrap();
let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
let corrupt_documents = [
br#"{"#.to_vec(),
br#"{"profile":"everyday","custom_max_model_mb":null,"extra":true}"#.to_vec(),
br#"{"profile":"custom","custom_max_model_mb":null}"#.to_vec(),
br#"{"profile":"custom","custom_max_model_mb":1}"#.to_vec(),
br#"{"profile":"custom","custom_max_model_mb":513}"#.to_vec(),
br#"{"profile":"everyday","custom_max_model_mb":512}"#.to_vec(),
br#"{"profile":"local_focused","custom_max_model_mb":512}"#.to_vec(),
vec![b' '; MAX_POLICY_BYTES as usize + 1],
];
for document in corrupt_documents {
std::fs::write(repository.path(), &document).unwrap();
let recovered = repository.load_with_evidence().unwrap();
assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
assert!(recovered.warning.is_some());
assert_eq!(std::fs::read(repository.path()).unwrap(), document);
}
}
#[cfg(unix)]
#[test]
fn non_regular_policy_source_is_reported_as_corrupt_without_following_it() {
use std::os::unix::fs::symlink;
let dir = tempfile::tempdir().unwrap();
let repository = FileResourcePolicyRepository::new(dir.path().to_path_buf());
let target = dir.path().join("target.json");
std::fs::write(
&target,
br#"{"profile":"local_focused","custom_max_model_mb":null}"#,
)
.unwrap();
symlink(&target, repository.path()).unwrap();
let recovered = repository.load_with_evidence().unwrap();
assert_eq!(recovered.source, ResourcePolicyLoadSource::CorruptDefault);
assert!(recovered.warning.is_some());
assert!(repository.path().is_symlink());
}
#[test]
fn concurrent_saves_leave_one_complete_document_and_no_staging_files() {
let dir = tempfile::tempdir().unwrap();
let repository =
std::sync::Arc::new(FileResourcePolicyRepository::new(dir.path().to_path_buf()));
let mut writers = Vec::new();
for index in 0..16_u64 {
let repository = repository.clone();
writers.push(std::thread::spawn(move || {
repository
.save(&ResourcePolicy {
profile: ResourceProfile::Custom,
custom_max_model_mb: Some(index * 512),
})
.unwrap();
}));
}
for writer in writers {
writer.join().unwrap();
}
let loaded = repository.load().unwrap();
assert_eq!(loaded.profile, ResourceProfile::Custom);
assert!(loaded.custom_max_model_mb.unwrap().is_multiple_of(512));
assert_no_temp_files(dir.path());
}
#[test]
fn cuda_uses_separate_vram_fit_and_host_ram_policy() {
let hardware = hardware(
64 * 1024,
crate::hardware::GpuBackend::Cuda,
Some(12 * 1024),
);
let evidence = evaluate_resources(&hardware, &ResourcePolicy::everyday());
assert_eq!(evidence.host_memory.configured_model_ceiling_mb, 26_214);
let accelerator = evidence.accelerator_memory.unwrap();
assert_eq!(accelerator.total_mb, 12 * 1024);
assert_eq!(accelerator.budget_mb, 12 * 1024);
assert_ne!(
accelerator.budget_mb,
evidence.host_memory.configured_model_ceiling_mb
);
}
#[test]
fn resource_preflight_zero_custom_budget_blocks_new_load() {
let coordinator = LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(0.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let preflight = coordinator.preflight(&model, 2_048);
assert_eq!(preflight.verdict, LocalLoadVerdict::DisabledByPolicy);
}
#[test]
fn resource_preflight_zero_budget_does_not_kill_resident_inference() {
let coordinator = LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(0.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
coordinator.mark_resident(&model.id, 2_400);
let preflight = coordinator.preflight(&model, 2_048);
assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
assert_eq!(
preflight.estimated_incremental_mb,
preflight.estimate.context_overhead_mb
+ preflight.estimate.runtime_overhead_mb
+ preflight.estimate.transient_margin_mb
);
}
#[test]
fn resource_preflight_resident_weights_are_incremental_only() {
let coordinator = LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(6_000)),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
coordinator.mark_resident(&model.id, 2_400);
let preflight = coordinator.preflight(&model, 2_048);
assert_eq!(
preflight.estimated_incremental_mb,
preflight.estimate.context_overhead_mb
+ preflight.estimate.runtime_overhead_mb
+ preflight.estimate.transient_margin_mb
);
assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
}
#[test]
fn resource_preflight_unavailable_live_probe_is_explicit() {
let coordinator = LocalAdmissionCoordinator::with_probe(
ResourcePolicy::everyday(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::unknown()),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let preflight = coordinator.preflight(&model, 2_048);
assert_eq!(preflight.live_available_mb, None);
assert_eq!(preflight.verdict, LocalLoadVerdict::LiveMemoryUnknown);
}
#[test]
fn resource_preflight_simultaneous_reservations_are_atomic() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(4.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(6_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
coordinator.mark_resident(&model.id, 2_400);
let first = coordinator.reserve(&model, 2_048).unwrap();
let second = coordinator.reserve(&model, 2_048).unwrap_err();
assert_eq!(
second.preflight.verdict,
LocalLoadVerdict::ExceedsConfiguredCeiling
);
drop(first);
assert!(coordinator.reserve(&model, 2_048).is_ok());
}
#[test]
fn distinct_state_roots_reserve_atomically_against_one_machine_ledger() {
let machine_ledger = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
let policy = ResourcePolicy::custom_gb(6.0).unwrap();
let hardware = hardware(32 * 1024, GpuBackend::Metal, None);
let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
policy.clone(),
hardware.clone(),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
machine_ledger.clone(),
));
let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
policy,
hardware,
Arc::new(FixedLiveMemoryProbe::known(24_000)),
machine_ledger,
));
let mut model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
model.cost.ram_mb = Some(2 * 1024);
model.cost.size_mb = Some(2 * 1024);
let barrier = Arc::new(std::sync::Barrier::new(2));
let attempts = [first, second].map(|coordinator| {
let model = model.clone();
let barrier = barrier.clone();
std::thread::spawn(move || {
barrier.wait();
coordinator.reserve(&model, 0)
})
});
let outcomes = attempts.map(|attempt| attempt.join().unwrap());
assert_eq!(outcomes.iter().filter(|result| result.is_ok()).count(), 1);
assert_eq!(
outcomes
.iter()
.find_map(|result| result.as_ref().err())
.expect("one cross-root request must be blocked")
.preflight
.verdict,
LocalLoadVerdict::ExceedsConfiguredCeiling
);
}
#[test]
fn resource_preflight_cuda_charges_weights_to_vram_and_overhead_to_host() {
let coordinator = LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(2.0).unwrap(),
hardware(64 * 1024, GpuBackend::Cuda, Some(12 * 1024)),
Arc::new(FixedLiveMemoryProbe::known(20_000)),
);
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let preflight = coordinator.preflight(&model, 2_048);
assert_eq!(preflight.resident_model_mb, 0);
assert_eq!(preflight.accelerator_total_mb, Some(12 * 1024));
assert_eq!(
preflight.accelerator_incremental_mb,
Some(preflight.estimate.weights_mb)
);
assert_eq!(
preflight.estimated_incremental_mb,
preflight.estimate.context_overhead_mb
+ preflight.estimate.runtime_overhead_mb
+ preflight.estimate.transient_margin_mb
);
assert_eq!(preflight.verdict, LocalLoadVerdict::Allowed);
}
#[test]
fn resource_preflight_cuda_vram_reservations_are_atomic() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(64 * 1024, GpuBackend::Cuda, Some(3_000)),
Arc::new(FixedLiveMemoryProbe::known(20_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let first = coordinator.reserve(&model, 2_048).unwrap();
assert_eq!(
coordinator
.reserve(&model, 2_048)
.unwrap_err()
.preflight
.verdict,
LocalLoadVerdict::ExceedsConfiguredCeiling
);
drop(first);
assert!(coordinator.reserve(&model, 2_048).is_ok());
}
#[test]
fn resource_preflight_model_maintenance_races_atomically_with_reserve() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::everyday(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let active = coordinator.reserve(&model, 2_048).unwrap();
assert!(matches!(
coordinator.begin_model_maintenance(&model.id),
Err(ModelMaintenanceError::ModelInUse(_))
));
drop(active);
let maintenance = coordinator.begin_model_maintenance(&model.id).unwrap();
assert_eq!(
coordinator
.reserve(&model, 2_048)
.unwrap_err()
.preflight
.verdict,
LocalLoadVerdict::ModelMaintenance
);
drop(maintenance);
assert!(coordinator.reserve(&model, 2_048).is_ok());
}
#[test]
fn measured_weights_are_rechecked_atomically_before_allocation() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(4.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(16_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
let blocked = reservation
.reconcile_measured_weights(6 * 1024 * 1024 * 1024)
.unwrap_err();
assert_eq!(
blocked.preflight.verdict,
LocalLoadVerdict::ExceedsConfiguredCeiling
);
assert_eq!(
coordinator.preflight(&model, 2_048).active_reservations_mb,
reservation.reserved_incremental_mb(),
"a rejected resize must not mutate the live reservation"
);
}
#[test]
fn cache_publication_transfers_cold_reservation_without_double_counting() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(16_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
let measured_bytes = 3 * 1024 * 1024 * 1024_u64;
reservation
.reconcile_measured_weights(measured_bytes)
.unwrap();
reservation.publish_resident_weights(measured_bytes);
let after = coordinator.preflight(&model, 2_048);
let request_overhead = after.estimate.context_overhead_mb
+ after.estimate.runtime_overhead_mb
+ after.estimate.transient_margin_mb;
assert_eq!(after.resident_model_mb, 3 * 1024);
assert_eq!(after.active_reservations_mb, request_overhead);
assert_eq!(reservation.reserved_incremental_mb(), request_overhead);
}
#[test]
fn resident_publication_never_reprobes_or_rejects_after_allocation() {
struct CountingProbe {
calls: AtomicU64,
available_mb: u64,
}
impl LiveMemoryProbe for CountingProbe {
fn available_memory_mb(&self) -> Result<Option<u64>, ResourcePolicyError> {
self.calls.fetch_add(1, Ordering::Relaxed);
Ok(Some(self.available_mb))
}
}
let probe = Arc::new(CountingProbe {
calls: AtomicU64::new(0),
available_mb: 24_000,
});
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
probe.clone(),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let mut reservation = coordinator.reserve(&model, 2_048).unwrap();
let measured = 3 * 1024 * 1024 * 1024_u64;
reservation.reconcile_measured_weights(measured).unwrap();
let calls_before_publish = probe.calls.load(Ordering::Relaxed);
reservation.publish_resident_weights(measured + 512 * 1024 * 1024);
assert_eq!(probe.calls.load(Ordering::Relaxed), calls_before_publish);
assert!(coordinator.is_resident(&model.id));
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.resident_models
.get(&model.id)
.map(|allocation| allocation.weights_mb),
Some(3 * 1024 + 512)
);
}
#[test]
fn simultaneous_cold_reservations_converge_on_one_resident_allocation() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(16.0).unwrap(),
hardware(32 * 1024, GpuBackend::Metal, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let mut first = coordinator.reserve(&model, 2_048).unwrap();
let mut second = coordinator.reserve(&model, 2_048).unwrap();
let measured = 3 * 1024 * 1024 * 1024_u64;
first.reconcile_measured_weights(measured).unwrap();
second.reconcile_measured_weights(measured).unwrap();
first.publish_resident_weights(measured);
second.publish_resident_weights(measured);
let after = coordinator.preflight(&model, 2_048);
let per_request_overhead = after.estimate.context_overhead_mb
+ after.estimate.runtime_overhead_mb
+ after.estimate.transient_margin_mb;
assert_eq!(after.resident_model_mb, 3 * 1024);
assert_eq!(
after.active_reservations_mb,
per_request_overhead * 2,
"peer publication must release only this reservation's redundant cold weights"
);
}
#[test]
fn resident_eviction_before_allocation_repromotes_reservation_to_cold() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let model = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.id == "mlx/qwen3-4b:4bit")
.unwrap();
let measured = 3 * 1024 * 1024 * 1024_u64;
let mut first = coordinator.reserve(&model, 2_048).unwrap();
first.reconcile_measured_weights(measured).unwrap();
first.publish_resident_weights(measured);
drop(first);
let mut replacement = coordinator.reserve(&model, 2_048).unwrap();
coordinator.mark_evicted(&model.id);
let preflight = replacement.reconcile_measured_weights(measured).unwrap();
assert!(preflight.estimated_incremental_mb >= 3 * 1024);
}
#[test]
fn maintenance_catalog_alias_blocks_provider_alias_reservation() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
coordinator
.register_model_aliases("mlx/kokoro-82m:6bit", ["mlx-community/Kokoro-82M-6bit"]);
let _maintenance = coordinator
.begin_model_maintenance("mlx/kokoro-82m:6bit")
.unwrap();
let mut provider_schema = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.is_local())
.expect("one local schema");
provider_schema.id = "mlx-community/Kokoro-82M-6bit".into();
let blocked = coordinator.reserve(&provider_schema, 512).unwrap_err();
assert_eq!(
blocked.preflight.verdict,
LocalLoadVerdict::ModelMaintenance
);
assert_eq!(blocked.preflight.model_id, "mlx/kokoro-82m:6bit");
}
#[test]
fn late_alias_registration_rekeys_active_and_resident_state_atomically() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let mut schema = crate::registry::builtin_catalog()
.into_iter()
.find(|model| model.is_local())
.unwrap();
schema.id = "provider/artifact".into();
let active = coordinator.reserve(&schema, 0).unwrap();
coordinator.register_model_aliases("catalog/model:default", ["provider/artifact"]);
assert!(matches!(
coordinator.begin_model_maintenance("catalog/model:default"),
Err(ModelMaintenanceError::ModelInUse(_))
));
drop(active);
let maintenance = coordinator
.begin_model_maintenance("catalog/model:default")
.unwrap();
let error = coordinator.reserve(&schema, 0).unwrap_err();
assert_eq!(error.preflight.verdict, LocalLoadVerdict::ModelMaintenance);
drop(maintenance);
}
#[test]
fn process_allocations_teardown_by_exact_owner_without_sibling_erasure() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(16.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let logical = "local/shared-model";
let worker = worker_process_allocation_id(logical);
let vllm = vllm_process_allocation_id(logical);
let bytes = 512 * 1024 * 1024_u64;
let mut worker_load = coordinator
.reserve_measured_host_allocation(logical, &worker, bytes, 0)
.unwrap();
worker_load.publish_resident_weights_as(&worker, bytes);
drop(worker_load);
let mut vllm_load = coordinator
.reserve_measured_host_allocation(logical, &vllm, bytes, 0)
.unwrap();
vllm_load.publish_resident_weights_as(&vllm, bytes);
drop(vllm_load);
assert_eq!(coordinator.resident_model_mb(), 1024);
coordinator.mark_teardown_pending_allocation(logical, &worker);
coordinator.mark_teardown_pending_allocation(logical, &vllm);
coordinator.finish_teardown_allocation(logical, &worker);
assert!(coordinator.teardown_pending(logical));
assert_eq!(
coordinator.resident_allocation_ids(logical),
vec![vllm.clone()]
);
assert_eq!(coordinator.resident_model_mb(), 512);
coordinator.finish_teardown_allocation(logical, &vllm);
assert!(!coordinator.teardown_pending(logical));
assert!(!coordinator.is_resident(logical));
}
#[test]
fn replacement_generation_reconcile_removes_peer_resident_discount_before_load() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(1.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let logical = "local/replaced-worker";
let bytes = 512 * 1024 * 1024_u64;
let mut old = coordinator
.reserve_measured_host_allocation(logical, "worker:old", bytes, 0)
.unwrap();
old.publish_resident_weights_as("worker:old", bytes);
drop(old);
let mut replacement = coordinator
.reserve_measured_host(logical, bytes, 0)
.expect("logical preflight initially sees the old resident");
assert_eq!(replacement.reserved_incremental_mb(), 0);
replacement.bind_allocation_id("worker:new");
replacement
.reconcile_measured_weights(bytes)
.expect("the exact replacement generation fits by itself");
assert_eq!(replacement.reserved_incremental_mb(), 512);
assert!(
coordinator
.reserve_measured_host("different/model", bytes, 0)
.is_err(),
"a second cold model must see both old residency and the replacement generation"
);
drop(replacement);
assert!(coordinator
.reserve_measured_host("different/model", bytes, 0)
.is_ok());
}
#[test]
fn pre_ack_pending_charge_blocks_other_scope_until_exact_exit_ack() {
let machine = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
ResourcePolicy::custom_gb(2.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
machine.clone(),
));
let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
ResourcePolicy::custom_gb(2.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
machine,
));
let logical = "managed/starting";
let allocation = worker_process_allocation_id(logical);
let cold = first
.reserve_measured_host_allocation(logical, &allocation, 1024 * 1024 * 1024, 0)
.unwrap();
first.mark_teardown_pending_allocation_with_charge(
logical,
&allocation,
cold.reconciled_weights_bytes(),
);
drop(cold);
let blocked = second
.reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
.unwrap_err();
assert_eq!(
blocked.preflight.verdict,
LocalLoadVerdict::ExceedsConfiguredCeiling
);
first.finish_teardown_allocation(logical, "worker:unrelated-sibling");
assert!(first.teardown_pending(logical));
assert_eq!(first.resident_model_mb(), 1024);
first.finish_teardown_allocation(logical, &allocation);
assert!(!first.teardown_pending(logical));
assert!(second
.reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
.is_ok());
}
#[test]
fn detached_native_lease_keeps_machine_charge_after_request_cancellation() {
let machine = Arc::new(Mutex::new(MachineAdmissionLedger::default()));
let first = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
ResourcePolicy::custom_gb(2.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
machine.clone(),
));
let second = Arc::new(LocalAdmissionCoordinator::with_probe_and_ledger(
ResourcePolicy::custom_gb(2.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
machine,
));
let reservation = first
.reserve_measured_host("detached/model", 1024 * 1024 * 1024, 0)
.unwrap();
let detached = reservation.detached_lease();
drop(reservation);
assert!(second
.reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
.is_err());
drop(detached);
assert!(second
.reserve_measured_host("other/model", 1536 * 1024 * 1024, 0)
.is_ok());
}
#[test]
fn normal_awaited_detached_work_shares_one_charge_with_its_request() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(1.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let request = coordinator
.reserve_measured_host("native/model-a", 512 * 1024 * 1024, 0)
.unwrap();
let detached = request.detached_lease();
let peer = coordinator
.reserve_measured_host("native/model-b", 512 * 1024 * 1024, 0)
.expect("request + its detached job are one 512 MB allocation, not two");
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_host_reservations_mb,
1024
);
drop(peer);
drop(detached);
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_host_reservations_mb,
512
);
drop(request);
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_host_reservations_mb,
0
);
}
#[test]
fn concurrent_non_aligned_starts_transfer_cold_weights_to_pending_without_double_charge() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
let first_allocation = vllm_process_allocation_id("managed/model-a");
let second_allocation = vllm_process_allocation_id("managed/model-b");
let mib = 1024 * 1024_u64;
let first_measured_bytes = 1024 * mib + 1;
let second_measured_bytes = 512 * mib + 1;
let mut first = coordinator
.reserve_measured_host_allocation(
"managed/model-a",
&first_allocation,
first_measured_bytes,
128,
)
.unwrap();
let mut second = coordinator
.reserve_measured_host_allocation(
"managed/model-b",
&second_allocation,
second_measured_bytes,
256,
)
.unwrap();
first.transfer_cold_weights_to_pending_allocation(&first_allocation, first_measured_bytes);
second
.transfer_cold_weights_to_pending_allocation(&second_allocation, second_measured_bytes);
assert_eq!(
coordinator.resident_model_mb(),
1538,
"each non-MiB-aligned allocation must be rounded up independently"
);
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_host_reservations_mb,
384,
"pending weights replace cold request weights while request overhead remains active"
);
first.publish_resident_weights_as(&first_allocation, first_measured_bytes);
second.publish_resident_weights_as(&second_allocation, second_measured_bytes);
assert_eq!(coordinator.resident_model_mb(), 1538);
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_host_reservations_mb,
384,
"publication must not subtract transferred weights twice"
);
drop(first);
drop(second);
assert_eq!(
coordinator
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.active_host_reservations_mb,
0
);
assert_eq!(coordinator.resident_model_mb(), 1538);
}
#[test]
fn pending_allocation_identity_is_idempotent_and_sibling_exact() {
let coordinator = Arc::new(LocalAdmissionCoordinator::with_probe(
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
Arc::new(FixedLiveMemoryProbe::known(24_000)),
));
coordinator.mark_teardown_pending_allocation_with_charge(
"same/model",
"worker:same/model",
512 * 1024 * 1024,
);
coordinator.mark_teardown_pending_allocation_with_charge(
"same/model",
"worker:same/model",
512 * 1024 * 1024,
);
coordinator.mark_teardown_pending_allocation_with_charge(
"same/model",
"vllm:same/model",
256 * 1024 * 1024,
);
assert_eq!(coordinator.resident_model_mb(), 768);
coordinator.finish_teardown_allocation("same/model", "worker:same/model");
assert!(coordinator.teardown_pending("same/model"));
assert_eq!(coordinator.resident_model_mb(), 256);
coordinator.finish_teardown_allocation("same/model", "worker:same/model");
assert_eq!(coordinator.resident_model_mb(), 256);
coordinator.finish_teardown_allocation("same/model", "vllm:same/model");
assert!(!coordinator.teardown_pending("same/model"));
}
#[test]
fn scoped_coordinator_identity_is_stable_for_one_state_root() {
let root = tempfile::tempdir().unwrap();
let first = scoped_local_admission(
root.path(),
ResourcePolicy::custom_gb(4.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
);
let second = scoped_local_admission(
root.path(),
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
);
assert!(Arc::ptr_eq(&first, &second));
assert_eq!(second.policy(), ResourcePolicy::custom_gb(8.0).unwrap());
}
#[cfg(unix)]
#[test]
fn scoped_coordinator_identity_unifies_symlinked_state_roots() {
use std::os::unix::fs::symlink;
let fixture = tempfile::tempdir().unwrap();
let real = fixture.path().join("real-state");
std::fs::create_dir(&real).unwrap();
let alias = fixture.path().join("state-alias");
symlink(&real, &alias).unwrap();
let first = scoped_local_admission(
&real,
ResourcePolicy::custom_gb(4.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
);
let second = scoped_local_admission(
&alias,
ResourcePolicy::custom_gb(8.0).unwrap(),
hardware(32 * 1024, GpuBackend::Cpu, None),
);
assert!(Arc::ptr_eq(&first, &second));
assert_eq!(second.policy(), ResourcePolicy::custom_gb(8.0).unwrap());
}
#[test]
fn unavailable_state_root_normalization_is_absolute_and_lexically_stable() {
let fixture = tempfile::tempdir().unwrap();
let missing = fixture.path().join("not-created").join("..").join("state");
assert_eq!(
normalized_state_root_key(&missing),
normalized_state_root_key(&fixture.path().join("state"))
);
}
#[cfg(unix)]
#[test]
fn missing_leaf_under_symlinked_parent_keeps_one_scope_identity() {
use std::os::unix::fs::symlink;
let fixture = tempfile::tempdir().unwrap();
let real = fixture.path().join("real");
std::fs::create_dir(&real).unwrap();
let alias = fixture.path().join("alias");
symlink(&real, &alias).unwrap();
assert_eq!(
normalized_state_root_key(&alias.join("missing").join("state")),
normalized_state_root_key(&real.join("missing").join("state"))
);
}
#[cfg(unix)]
#[test]
fn parent_components_are_resolved_after_symlinks_not_lexically_before_them() {
use std::os::unix::fs::symlink;
let fixture = tempfile::tempdir().unwrap();
let physical_parent = fixture.path().join("physical");
let physical_child = physical_parent.join("child");
std::fs::create_dir_all(&physical_child).unwrap();
let aliases = fixture.path().join("aliases");
std::fs::create_dir(&aliases).unwrap();
let alias = aliases.join("runtime");
symlink(&physical_child, &alias).unwrap();
let through_alias = alias.join("..").join("missing-state");
assert_eq!(
normalized_state_root_key(&through_alias),
normalized_state_root_key(&physical_parent.join("missing-state"))
);
assert_ne!(
normalized_state_root_key(&through_alias),
aliases.join("missing-state")
);
}
fn assert_no_temp_files(directory: &std::path::Path) {
let entries = std::fs::read_dir(directory)
.unwrap()
.map(|entry| entry.unwrap().path())
.collect::<Vec<_>>();
assert_eq!(entries, vec![directory.join("model-resource-policy.json")]);
}
#[cfg(unix)]
fn assert_private_mode(path: &std::path::Path, expected: u32) {
use std::os::unix::fs::PermissionsExt;
assert_eq!(
std::fs::metadata(path).unwrap().permissions().mode() & 0o777,
expected
);
}
#[cfg(not(unix))]
fn assert_private_mode(_path: &std::path::Path, _expected: u32) {}
}