use std::collections::HashMap;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use car_inference::model_management::{ModelLease, ModelManagementStore};
use car_inference::LocalAdmissionCoordinator;
use crate::VoiceError;
#[derive(Clone)]
pub(crate) struct VoiceLocalAdmission {
coordinator: Arc<LocalAdmissionCoordinator>,
management: Option<ModelManagementStore>,
state: Arc<VoiceAdmissionState>,
}
#[derive(Default)]
struct VoiceAdmissionState {
next_allocation: AtomicU64,
load_gates: Mutex<HashMap<String, Arc<Mutex<()>>>>,
}
impl VoiceLocalAdmission {
pub(crate) fn shared() -> Self {
static STATE: OnceLock<Arc<VoiceAdmissionState>> = OnceLock::new();
Self {
coordinator: car_inference::resource_policy::shared_local_admission(),
management: Some(ModelManagementStore::new(
car_home::root_or_relative(),
car_inference::default_models_dir(),
)),
state: STATE
.get_or_init(|| Arc::new(VoiceAdmissionState::default()))
.clone(),
}
}
#[cfg(test)]
pub(crate) fn new(coordinator: Arc<LocalAdmissionCoordinator>) -> Self {
Self {
coordinator,
management: None,
state: Arc::new(VoiceAdmissionState::default()),
}
}
#[cfg(test)]
pub(crate) fn new_with_management(
coordinator: Arc<LocalAdmissionCoordinator>,
management: ModelManagementStore,
) -> Self {
Self {
coordinator,
management: Some(management),
state: Arc::new(VoiceAdmissionState::default()),
}
}
pub(crate) fn load_installed<T>(
&self,
model_id: &str,
installed_path: &Path,
request_overhead_mb: u64,
loader: impl FnOnce() -> Result<T, VoiceError>,
) -> Result<(T, VoiceResidentLease), VoiceError> {
let canonical_model_id = catalog_canonical_id(model_id);
self.coordinator.register_model_aliases(
&canonical_model_id,
catalog_aliases(model_id, Some(installed_path)),
);
let measured_bytes = car_inference::backend_cache::estimate_model_size(installed_path);
self.load_measured(
&canonical_model_id,
measured_bytes,
request_overhead_mb,
loader,
)
}
pub(crate) fn load_declared<T>(
&self,
model_id: &str,
measured_bytes: u64,
request_overhead_mb: u64,
loader: impl FnOnce() -> Result<T, VoiceError>,
) -> Result<(T, VoiceResidentLease), VoiceError> {
let canonical_model_id = catalog_canonical_id(model_id);
self.coordinator
.register_model_aliases(&canonical_model_id, catalog_aliases(model_id, None));
self.load_measured(
&canonical_model_id,
measured_bytes,
request_overhead_mb,
loader,
)
}
fn load_measured<T>(
&self,
model_id: &str,
measured_bytes: u64,
request_overhead_mb: u64,
loader: impl FnOnce() -> Result<T, VoiceError>,
) -> Result<(T, VoiceResidentLease), VoiceError> {
let load_gate = {
let mut gates = self
.state
.load_gates
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
Arc::clone(
gates
.entry(model_id.to_string())
.or_insert_with(|| Arc::new(Mutex::new(()))),
)
};
let singleflight = load_gate
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let result = (|| {
let activity_lease = self
.management
.as_ref()
.map(|management| management.acquire_lease(model_id))
.transpose()
.map_err(|error| VoiceError::Config(error.to_string()))?;
let allocation = self.state.next_allocation.fetch_add(1, Ordering::AcqRel);
let allocation_id = format!("{model_id}#voice-allocation-{allocation}");
let mut reservation = self
.coordinator
.reserve_measured_host_allocation(
model_id,
&allocation_id,
measured_bytes,
request_overhead_mb,
)
.map_err(|error| VoiceError::Config(error.to_string()))?;
let loaded = loader()?;
reservation.publish_resident_weights(measured_bytes);
Ok((
loaded,
VoiceResidentLease {
inner: Arc::new(VoiceResidentLeaseInner {
model_id: allocation_id,
coordinator: self.coordinator.clone(),
_activity_lease: activity_lease,
}),
},
))
})();
drop(singleflight);
let mut gates = self
.state
.load_gates
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if Arc::strong_count(&load_gate) == 2 {
gates.remove(model_id);
}
result
}
}
fn catalog_canonical_id(model_id: &str) -> String {
let lower = model_id.to_ascii_lowercase();
if lower.contains("parakeet-tdt-0.6b") {
return "mlx/parakeet-tdt-0.6b-v3:default".to_string();
}
if lower.contains("kokoro-82m") {
if lower.contains("6bit") {
return "mlx/kokoro-82m:6bit".to_string();
}
if lower.contains("bf16") {
return "mlx/kokoro-82m:bf16".to_string();
}
}
if let Some(variant) = model_id
.strip_prefix("voice/whisper-cpp/")
.or_else(|| model_id.strip_prefix("whisper/"))
{
let variant = variant.strip_suffix(":default").unwrap_or(variant);
return format!("whisper/{variant}:default");
}
model_id.to_string()
}
pub(crate) fn catalog_aliases(model_id: &str, installed_path: Option<&Path>) -> Vec<String> {
let mut aliases = vec![model_id.to_string()];
if let Some(alias) = model_id.strip_prefix("voice/") {
aliases.push(alias.to_string());
}
if let Some(alias) = model_id.rsplit('/').next() {
aliases.push(alias.to_string());
}
let lower = model_id.to_ascii_lowercase();
if lower.contains("parakeet-tdt-0.6b") {
aliases.extend([
"mlx/parakeet-tdt-0.6b-v3:default".to_string(),
"mlx-community/parakeet-tdt-0.6b-v3".to_string(),
"parakeet-tdt-0.6b-v2-int8".to_string(),
]);
}
if let Some(variant) = model_id
.strip_prefix("voice/whisper-cpp/")
.or_else(|| model_id.strip_prefix("whisper/"))
{
let variant = variant.strip_suffix(":default").unwrap_or(variant);
aliases.push(format!("whisper/{variant}:default"));
aliases.push(format!("ggml-{variant}.bin"));
}
if let Some(path) = installed_path {
if let Some(alias) = path.file_name().and_then(|value| value.to_str()) {
aliases.push(alias.to_string());
}
}
aliases.sort();
aliases.dedup();
aliases
}
#[derive(Clone)]
pub(crate) struct VoiceResidentLease {
inner: Arc<VoiceResidentLeaseInner>,
}
struct VoiceResidentLeaseInner {
model_id: String,
coordinator: Arc<LocalAdmissionCoordinator>,
_activity_lease: Option<ModelLease>,
}
impl std::fmt::Debug for VoiceResidentLease {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("VoiceResidentLease")
.field("model_id", &self.inner.model_id)
.finish_non_exhaustive()
}
}
impl Drop for VoiceResidentLeaseInner {
fn drop(&mut self) {
self.coordinator.mark_evicted(&self.model_id);
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use car_inference::hardware::{GpuBackend, HardwareInfo};
use car_inference::model_management::ModelManagementStore;
use car_inference::{LocalAdmissionCoordinator, ResourcePolicy};
use super::{catalog_aliases, catalog_canonical_id, VoiceLocalAdmission};
#[test]
fn native_voice_ids_use_exact_catalog_management_identities() {
for (native_id, catalog_id) in [
(
"voice/parakeet-tdt-0.6b",
"mlx/parakeet-tdt-0.6b-v3:default",
),
(
"voice/whisper-cpp/large-v3-turbo-q5_0",
"whisper/large-v3-turbo-q5_0:default",
),
("mlx-community/Kokoro-82M-6bit", "mlx/kokoro-82m:6bit"),
] {
assert_eq!(catalog_canonical_id(native_id), catalog_id);
}
assert_eq!(catalog_canonical_id("voice/custom"), "voice/custom");
}
#[test]
fn coordinator_canonicalizes_kokoro_catalog_ids_without_voice_alias_duplication() {
let coordinator = LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo::detect(),
);
for (artifact, catalog_id) in [
("mlx-community/Kokoro-82M-6bit", "mlx/kokoro-82m:6bit"),
("mlx-community/Kokoro-82M-bf16", "mlx/kokoro-82m:bf16"),
] {
assert_eq!(coordinator.canonical_model_id(artifact), catalog_id);
}
}
#[test]
fn voice_download_remains_but_native_loader_is_blocked_by_shared_policy() {
let dir = tempfile::tempdir().unwrap();
let weights = dir.path().join("voice-model.bin");
std::fs::write(&weights, vec![0_u8; 1024 * 1024]).unwrap();
let admission = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(0.0).unwrap(),
HardwareInfo {
os: "test".into(),
arch: "test".into(),
cpu_cores: 8,
total_ram_mb: 32 * 1024,
gpu_backend: GpuBackend::Cpu,
gpu_memory_mb: None,
gpu_devices: Vec::new(),
recommended_model: "fixture".into(),
recommended_context: 4096,
max_model_mb: 32 * 1024,
},
));
let gate = VoiceLocalAdmission::new(admission);
let loaded = AtomicBool::new(false);
let error = gate
.load_installed("voice/whisper", &weights, 128, || {
loaded.store(true, Ordering::SeqCst);
Ok::<_, crate::VoiceError>(())
})
.unwrap_err();
assert!(error.to_string().contains("blocked"));
assert!(!loaded.load(Ordering::SeqCst));
assert!(
weights.exists(),
"admission must never delete a downloaded model"
);
}
#[test]
fn duplicate_voice_allocations_are_fully_charged_and_drop_independently() {
let dir = tempfile::tempdir().unwrap();
let weights = dir.path().join("voice-model.bin");
std::fs::write(&weights, vec![0_u8; 1024 * 1024]).unwrap();
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let gate = VoiceLocalAdmission::new(coordinator.clone());
let loads = AtomicUsize::new(0);
let (_, first) = gate
.load_installed("voice/duplicate", &weights, 0, || {
loads.fetch_add(1, Ordering::SeqCst);
Ok::<_, crate::VoiceError>(())
})
.unwrap();
let (_, second) = gate
.load_installed("voice/duplicate", &weights, 0, || {
loads.fetch_add(1, Ordering::SeqCst);
Ok::<_, crate::VoiceError>(())
})
.unwrap();
assert_eq!(loads.load(Ordering::SeqCst), 2);
assert_eq!(coordinator.resident_model_mb(), 2);
assert!(coordinator.is_resident("voice/duplicate"));
drop(first);
assert_eq!(coordinator.resident_model_mb(), 1);
assert!(coordinator.is_resident("voice/duplicate"));
drop(second);
assert_eq!(coordinator.resident_model_mb(), 0);
assert!(!coordinator.is_resident("voice/duplicate"));
}
#[test]
fn whisper_variants_in_one_directory_keep_distinct_identities_and_reservations() {
let dir = tempfile::tempdir().unwrap();
let small = dir.path().join("ggml-small.en-q5_1.bin");
let large = dir.path().join("ggml-large-v3-turbo-q5_0.bin");
std::fs::write(&small, vec![0_u8; 1024 * 1024]).unwrap();
std::fs::write(&large, vec![0_u8; 1024 * 1024]).unwrap();
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let gate = VoiceLocalAdmission::new(coordinator.clone());
let (_, small_lease) = gate
.load_installed("voice/whisper-cpp/small.en-q5_1", &small, 0, || {
Ok::<_, crate::VoiceError>(())
})
.unwrap();
let (_, large_lease) = gate
.load_installed("voice/whisper-cpp/large-v3-turbo-q5_0", &large, 0, || {
Ok::<_, crate::VoiceError>(())
})
.unwrap();
assert_ne!(
coordinator.canonical_model_id("voice/whisper-cpp/small.en-q5_1"),
coordinator.canonical_model_id("voice/whisper-cpp/large-v3-turbo-q5_0")
);
assert_eq!(
coordinator
.resident_allocation_ids("voice/whisper-cpp/small.en-q5_1")
.len(),
1
);
assert_eq!(
coordinator
.resident_allocation_ids("voice/whisper-cpp/large-v3-turbo-q5_0")
.len(),
1
);
assert_eq!(coordinator.resident_model_mb(), 2);
drop(small_lease);
assert!(coordinator.is_resident("voice/whisper-cpp/large-v3-turbo-q5_0"));
drop(large_lease);
}
#[test]
fn artifact_alias_finds_voice_residency_and_final_lease_drop_clears_it() {
let dir = tempfile::tempdir().unwrap();
let artifact_dir = dir.path().join("parakeet-tdt-0.6b-v2-int8");
std::fs::create_dir(&artifact_dir).unwrap();
std::fs::write(artifact_dir.join("model.onnx"), vec![0_u8; 1024 * 1024]).unwrap();
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let (_, lease) = VoiceLocalAdmission::new(coordinator.clone())
.load_installed("voice/parakeet-tdt-0.6b", &artifact_dir, 0, || {
Ok::<_, crate::VoiceError>(())
})
.unwrap();
assert!(coordinator.is_resident("parakeet-tdt-0.6b-v2-int8"));
assert_eq!(
coordinator
.resident_allocation_ids("parakeet-tdt-0.6b-v2-int8")
.len(),
1
);
let detached_work = lease.clone();
drop(lease);
assert!(
coordinator.is_resident("parakeet-tdt-0.6b-v2-int8"),
"a detached provider operation must retain residency after the provider drops"
);
drop(detached_work);
assert!(!coordinator.is_resident("parakeet-tdt-0.6b-v2-int8"));
}
#[test]
fn resident_voice_load_holds_catalog_activity_lease_until_final_clone_drops() {
let dir = tempfile::tempdir().unwrap();
let models_dir = dir.path().join("models");
std::fs::create_dir(&models_dir).unwrap();
let artifact_dir = models_dir.join("parakeet-tdt-0.6b-v2-int8");
std::fs::create_dir(&artifact_dir).unwrap();
std::fs::write(artifact_dir.join("model.onnx"), vec![0_u8; 1024]).unwrap();
let loader_store = ModelManagementStore::new(dir.path().join("loader"), models_dir.clone());
let removal_store = ModelManagementStore::new(dir.path().join("remover"), models_dir);
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let gate = VoiceLocalAdmission::new_with_management(coordinator, loader_store);
let (_, resident) = gate
.load_installed("voice/parakeet-tdt-0.6b", &artifact_dir, 0, || {
Ok::<_, crate::VoiceError>(())
})
.unwrap();
let catalog_id = "mlx/parakeet-tdt-0.6b-v3:default";
assert!(
removal_store.model_in_use(catalog_id).unwrap(),
"an independent remover must observe the resident native voice allocation"
);
assert!(
!removal_store
.model_in_use("voice/parakeet-tdt-0.6b")
.unwrap(),
"the activity lock must use the removable catalog identity, not a provider alias"
);
let detached = resident.clone();
drop(resident);
assert!(
removal_store.model_in_use(catalog_id).unwrap(),
"a detached voice operation must retain the cross-process activity lease"
);
drop(detached);
assert!(
!removal_store.model_in_use(catalog_id).unwrap(),
"the exact activity lease must release after the final resident clone drops"
);
}
#[test]
fn real_catalog_ids_resolve_live_voice_allocations() {
let dir = tempfile::tempdir().unwrap();
let parakeet = dir.path().join("parakeet-tdt-0.6b-v2-int8");
std::fs::create_dir(¶keet).unwrap();
std::fs::write(parakeet.join("model.onnx"), vec![0_u8; 1024]).unwrap();
let whisper = dir.path().join("ggml-large-v3-turbo-q5_0.bin");
std::fs::write(&whisper, vec![0_u8; 1024]).unwrap();
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let gate = VoiceLocalAdmission::new(coordinator.clone());
let (_, parakeet_lease) = gate
.load_installed("voice/parakeet-tdt-0.6b", ¶keet, 0, || {
Ok::<_, crate::VoiceError>(())
})
.unwrap();
let (_, whisper_lease) = gate
.load_installed("voice/whisper-cpp/large-v3-turbo-q5_0", &whisper, 0, || {
Ok::<_, crate::VoiceError>(())
})
.unwrap();
assert!(!coordinator
.resident_allocation_ids("mlx/parakeet-tdt-0.6b-v3:default")
.is_empty());
assert!(!coordinator
.resident_allocation_ids("whisper/large-v3-turbo-q5_0:default")
.is_empty());
for (repo, catalog_id) in [
("mlx-community/Kokoro-82M-6bit", "mlx/kokoro-82m:6bit"),
("mlx-community/Kokoro-82M-bf16", "mlx/kokoro-82m:bf16"),
] {
coordinator.register_model_aliases(repo, catalog_aliases(repo, None));
coordinator.mark_resident(repo, 1);
assert!(
!coordinator.resident_allocation_ids(catalog_id).is_empty(),
"catalog removal id {catalog_id} must resolve the live Kokoro allocation"
);
coordinator.mark_evicted(repo);
}
drop(parakeet_lease);
drop(whisper_lease);
}
#[test]
fn duplicate_voice_loaders_are_singleflight_per_model() {
let dir = tempfile::tempdir().unwrap();
let weights = dir.path().join("voice-model.bin");
std::fs::write(&weights, vec![0_u8; 1024 * 1024]).unwrap();
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(8.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let gate = VoiceLocalAdmission::new(coordinator);
let active = Arc::new(AtomicUsize::new(0));
let peak = Arc::new(AtomicUsize::new(0));
let mut threads = Vec::new();
for _ in 0..2 {
let gate = gate.clone();
let weights = weights.clone();
let active = active.clone();
let peak = peak.clone();
threads.push(std::thread::spawn(move || {
gate.load_installed("voice/singleflight", &weights, 0, || {
let now = active.fetch_add(1, Ordering::SeqCst) + 1;
peak.fetch_max(now, Ordering::SeqCst);
std::thread::sleep(std::time::Duration::from_millis(25));
active.fetch_sub(1, Ordering::SeqCst);
Ok::<_, crate::VoiceError>(())
})
.unwrap()
}));
}
let leases = threads
.into_iter()
.map(|thread| thread.join().unwrap().1)
.collect::<Vec<_>>();
assert_eq!(peak.load(Ordering::SeqCst), 1);
drop(leases);
}
#[test]
fn onnx_vad_declared_allocation_obeys_zero_gb_policy() {
let coordinator = Arc::new(LocalAdmissionCoordinator::new(
ResourcePolicy::custom_gb(0.0).unwrap(),
HardwareInfo {
total_ram_mb: 32 * 1024,
..HardwareInfo::detect()
},
));
let loaded = AtomicBool::new(false);
let error = VoiceLocalAdmission::new(coordinator)
.load_declared("voice/silero-vad-v5-onnx", 2 * 1024 * 1024, 16, || {
loaded.store(true, Ordering::SeqCst);
Ok::<_, crate::VoiceError>(())
})
.unwrap_err();
assert!(error.to_string().contains("blocked"));
assert!(!loaded.load(Ordering::SeqCst));
}
}