use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{SystemTime, UNIX_EPOCH};
pub const DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB: u64 = 4096;
pub const DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB: u64 = 768;
pub const DEFAULT_IDLE_TTL_SECONDS: u64 = 0;
pub const DEFAULT_LLAMA_CPP_ENDPOINT: &str = "http://127.0.0.1:8080/models/unload";
pub const DEFAULT_OLLAMA_ENDPOINT: &str = "http://127.0.0.1:11434/api/generate";
pub const DEFAULT_VLLM_ENDPOINT: &str = "http://127.0.0.1:8000/sleep";
pub const LLAMA_CPP_ENDPOINT_ENV_VAR: &str = "TSIFT_LLAMA_CPP_ENDPOINT";
pub const OLLAMA_ENDPOINT_ENV_VAR: &str = "TSIFT_OLLAMA_ENDPOINT";
pub const VLLM_ENDPOINT_ENV_VAR: &str = "TSIFT_VLLM_ENDPOINT";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub enum ModelRole {
Extract,
Embed,
Rerank,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum ProviderKind {
LlamaCpp,
Ollama,
Vllm,
HashFallback,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum UnloadStrategy {
ProcessExit,
OllamaKeepAliveZero,
LlamaCppRouterUnload,
VllmSleep,
None,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum ConcurrencyClass {
ExclusiveLargeGpu,
SharedSmallGpu,
CpuOrHash,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum LeaseMode {
Exclusive,
Shared,
CpuOrHash,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum UnloadActionKind {
ProviderApi,
ProcessExit,
Sleep,
Noop,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ModelProfile {
pub id: &'static str,
pub label: &'static str,
pub provider: ProviderKind,
pub model_ref: &'static str,
pub quantization: &'static str,
pub roles: Vec<ModelRole>,
pub context_tokens: u32,
pub estimated_weights_mib: u64,
pub estimated_kv_mib: u64,
pub runtime_margin_mib: u64,
pub concurrency: ConcurrencyClass,
pub unload_strategy: UnloadStrategy,
pub notes: &'static str,
}
impl ModelProfile {
pub fn estimated_total_mib(&self) -> u64 {
self.estimated_weights_mib + self.estimated_kv_mib + self.runtime_margin_mib
}
pub fn supports_role(&self, role: &ModelRole) -> bool {
self.roles.iter().any(|candidate| candidate == role)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuProcess {
pub pid: Option<u32>,
pub process_name: String,
pub used_memory_mib: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuProbe {
pub timestamp_unix_seconds: Option<u64>,
pub available: bool,
pub gpu_name: Option<String>,
pub driver_version: Option<String>,
pub total_vram_mib: Option<u64>,
pub used_vram_mib: Option<u64>,
pub free_vram_mib: Option<u64>,
pub processes: Vec<GpuProcess>,
pub error: Option<String>,
}
impl GpuProbe {
pub fn unavailable(error: impl Into<String>) -> Self {
Self {
timestamp_unix_seconds: Some(current_unix_seconds()),
available: false,
gpu_name: None,
driver_version: None,
total_vram_mib: None,
used_vram_mib: None,
free_vram_mib: None,
processes: Vec::new(),
error: Some(error.into()),
}
}
pub fn synthetic_vram(used_vram_mib: u64) -> Self {
Self {
timestamp_unix_seconds: Some(current_unix_seconds()),
available: true,
gpu_name: Some("synthetic GPU".to_string()),
driver_version: None,
total_vram_mib: None,
used_vram_mib: Some(used_vram_mib),
free_vram_mib: None,
processes: Vec::new(),
error: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ProfileSelection {
pub profile: ModelProfile,
pub selectable: bool,
pub reason: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct LocalModelStatusReport {
pub gpu_probe: GpuProbe,
pub extractor_profiles: Vec<ProfileSelection>,
pub embedding_profiles: Vec<ProfileSelection>,
pub recommended_extractor: Option<String>,
pub recommended_embedding: Option<String>,
pub notes: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ProviderUnloadAction {
pub kind: UnloadActionKind,
pub label: String,
pub command: Option<Vec<String>>,
pub http_method: Option<String>,
pub endpoint: Option<String>,
pub body_json: Option<String>,
pub required: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct LocalModelLease {
pub lease_id: String,
pub mode: LeaseMode,
pub profile: ModelProfile,
pub pre_load_gpu_probe: GpuProbe,
pub provider_endpoint: Option<String>,
pub provider_pid: Option<u32>,
pub idle_ttl_seconds: u64,
pub unload_strategy: UnloadStrategy,
pub unload_actions: Vec<ProviderUnloadAction>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum VramCleanupStatus {
Proven,
ProvenByExternalAccounting,
NotProven,
ProbeUnavailable,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct VramCleanupEvaluation {
pub status: VramCleanupStatus,
pub cleanup_proven: bool,
pub pre_used_mib: Option<u64>,
pub post_used_mib: Option<u64>,
pub allowed_post_used_mib: Option<u64>,
pub used_delta_mib: Option<i64>,
pub external_process_delta_mib: u64,
pub blocking_processes: Vec<GpuProcess>,
pub reason: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct LocalModelLifecycleReport {
pub lease: LocalModelLease,
pub post_unload_gpu_probe: GpuProbe,
pub cleanup: VramCleanupEvaluation,
pub notes: Vec<String>,
}
pub fn default_model_profiles() -> Vec<ModelProfile> {
vec![
ModelProfile {
id: "qwen3-32b-q4",
label: "Qwen3-32B 4-bit",
provider: ProviderKind::LlamaCpp,
model_ref: "Qwen/Qwen3-32B-GGUF",
quantization: "q4",
roles: vec![ModelRole::Extract],
context_tokens: 32_768,
estimated_weights_mib: 20_500,
estimated_kv_mib: 4_096,
runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
concurrency: ConcurrencyClass::ExclusiveLargeGpu,
unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
notes: "default quality extractor/reasoner for a clear RTX 5090",
},
ModelProfile {
id: "qwen3-30b-a3b-instruct-2507-q4",
label: "Qwen3-30B-A3B-Instruct-2507 4-bit",
provider: ProviderKind::LlamaCpp,
model_ref: "Qwen/Qwen3-30B-A3B-Instruct-2507",
quantization: "q4",
roles: vec![ModelRole::Extract],
context_tokens: 262_144,
estimated_weights_mib: 19_000,
estimated_kv_mib: 4_096,
runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
concurrency: ConcurrencyClass::ExclusiveLargeGpu,
unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
notes: "throughput and long-context extractor fallback",
},
ModelProfile {
id: "qwen3-embedding-0.6b",
label: "Qwen3-Embedding-0.6B",
provider: ProviderKind::LlamaCpp,
model_ref: "Qwen/Qwen3-Embedding-0.6B-GGUF",
quantization: "q8_or_f16",
roles: vec![ModelRole::Embed, ModelRole::Rerank],
context_tokens: 32_768,
estimated_weights_mib: 1_200,
estimated_kv_mib: 512,
runtime_margin_mib: 1_024,
concurrency: ConcurrencyClass::SharedSmallGpu,
unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
notes: "default low-pressure embedding companion",
},
ModelProfile {
id: "qwen3-embedding-4b",
label: "Qwen3-Embedding-4B",
provider: ProviderKind::LlamaCpp,
model_ref: "Qwen/Qwen3-Embedding-4B",
quantization: "q4_or_q8",
roles: vec![ModelRole::Embed, ModelRole::Rerank],
context_tokens: 32_768,
estimated_weights_mib: 4_200,
estimated_kv_mib: 1_024,
runtime_margin_mib: 1_024,
concurrency: ConcurrencyClass::SharedSmallGpu,
unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
notes: "higher-quality embedding candidate",
},
ModelProfile {
id: "qwen3-embedding-8b",
label: "Qwen3-Embedding-8B",
provider: ProviderKind::LlamaCpp,
model_ref: "Qwen/Qwen3-Embedding-8B",
quantization: "q4_or_q8",
roles: vec![ModelRole::Embed, ModelRole::Rerank],
context_tokens: 32_768,
estimated_weights_mib: 8_200,
estimated_kv_mib: 2_048,
runtime_margin_mib: 2_048,
concurrency: ConcurrencyClass::SharedSmallGpu,
unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
notes: "benchmark when vector quality matters",
},
ModelProfile {
id: "qwen3.5-35b-a3b-q4",
label: "Qwen3.5-35B-A3B 4-bit",
provider: ProviderKind::LlamaCpp,
model_ref: "Qwen/Qwen3.5-35B-A3B",
quantization: "q4",
roles: vec![ModelRole::Extract],
context_tokens: 128_000,
estimated_weights_mib: 24_000,
estimated_kv_mib: 8_192,
runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
concurrency: ConcurrencyClass::ExclusiveLargeGpu,
unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
notes: "benchmark-only until a reduced-context single-5090 profile is proven",
},
ModelProfile {
id: "tsift-local-hash-v1",
label: "tsift local hash fallback",
provider: ProviderKind::HashFallback,
model_ref: "builtin",
quantization: "none",
roles: vec![ModelRole::Embed],
context_tokens: 0,
estimated_weights_mib: 0,
estimated_kv_mib: 0,
runtime_margin_mib: 0,
concurrency: ConcurrencyClass::CpuOrHash,
unload_strategy: UnloadStrategy::None,
notes: "deterministic fallback for tests and offline runs",
},
ModelProfile {
id: "qwen3-32b-q4-ollama",
label: "Qwen3-32B 4-bit (Ollama)",
provider: ProviderKind::Ollama,
model_ref: "hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M",
quantization: "q4",
roles: vec![ModelRole::Extract],
context_tokens: 32_768,
estimated_weights_mib: 20_500,
estimated_kv_mib: 4_096,
runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
concurrency: ConcurrencyClass::ExclusiveLargeGpu,
unload_strategy: UnloadStrategy::OllamaKeepAliveZero,
notes: "default Ollama-served quality extractor (lazy: keep_alive:0 unloads VRAM)",
},
ModelProfile {
id: "qwen3-embedding-0.6b-ollama",
label: "Qwen3-Embedding-0.6B (Ollama)",
provider: ProviderKind::Ollama,
model_ref: "hf.co/Qwen/Qwen3-Embedding-0.6B-GGUF",
quantization: "q8_or_f16",
roles: vec![ModelRole::Embed, ModelRole::Rerank],
context_tokens: 32_768,
estimated_weights_mib: 1_200,
estimated_kv_mib: 512,
runtime_margin_mib: 1_024,
concurrency: ConcurrencyClass::SharedSmallGpu,
unload_strategy: UnloadStrategy::OllamaKeepAliveZero,
notes: "default low-pressure Ollama-served embedding companion",
},
]
}
pub fn probe_nvidia_smi() -> GpuProbe {
let output = Command::new("nvidia-smi")
.args([
"--query-gpu=name,driver_version,memory.total,memory.used,memory.free",
"--format=csv,noheader,nounits",
])
.output();
let output = match output {
Ok(output) => output,
Err(error) => return GpuProbe::unavailable(format!("nvidia-smi unavailable: {error}")),
};
if !output.status.success() {
return GpuProbe::unavailable(format!(
"nvidia-smi failed: {}",
String::from_utf8_lossy(&output.stderr).trim()
));
}
let stdout = String::from_utf8_lossy(&output.stdout);
match parse_gpu_query(stdout.lines().next().unwrap_or_default()) {
Ok(mut probe) => {
probe.processes = query_nvidia_compute_processes();
probe
}
Err(error) => GpuProbe::unavailable(error.to_string()),
}
}
pub fn build_status_report(probe_gpu: bool) -> LocalModelStatusReport {
let gpu_probe = if probe_gpu {
probe_nvidia_smi()
} else {
GpuProbe::unavailable("gpu probe skipped")
};
build_status_report_with_probe(gpu_probe)
}
pub fn build_status_report_with_probe(gpu_probe: GpuProbe) -> LocalModelStatusReport {
let profiles = default_model_profiles();
let extractor_profiles = rank_profiles_for_role(&profiles, &gpu_probe, ModelRole::Extract);
let embedding_profiles = rank_profiles_for_role(&profiles, &gpu_probe, ModelRole::Embed);
let recommended_extractor = extractor_profiles
.iter()
.find(|selection| selection.selectable)
.map(|selection| selection.profile.id.to_string());
let recommended_embedding = embedding_profiles
.iter()
.find(|selection| selection.selectable)
.map(|selection| selection.profile.id.to_string());
let mut notes = vec![
"large 30B/32B extractor profiles are single-lease on one RTX 5090".to_string(),
"use provider unload hooks or process exit after each batch to clear VRAM".to_string(),
];
if !gpu_probe.available {
notes.push("GPU probe unavailable; profile fit is conservative".to_string());
}
LocalModelStatusReport {
gpu_probe,
extractor_profiles,
embedding_profiles,
recommended_extractor,
recommended_embedding,
notes,
}
}
pub fn profile_by_id(profile_id: &str) -> Option<ModelProfile> {
default_model_profiles()
.into_iter()
.find(|profile| profile.id == profile_id)
}
pub fn lease_mode_for_profile(profile: &ModelProfile) -> LeaseMode {
match profile.concurrency {
ConcurrencyClass::ExclusiveLargeGpu => LeaseMode::Exclusive,
ConcurrencyClass::SharedSmallGpu => LeaseMode::Shared,
ConcurrencyClass::CpuOrHash => LeaseMode::CpuOrHash,
}
}
pub fn build_local_model_lease(
profile: ModelProfile,
pre_load_gpu_probe: GpuProbe,
provider_endpoint: Option<String>,
provider_pid: Option<u32>,
idle_ttl_seconds: u64,
) -> LocalModelLease {
let timestamp = pre_load_gpu_probe
.timestamp_unix_seconds
.unwrap_or_else(current_unix_seconds);
let lease_id = format!("{}-{timestamp}", profile.id);
let unload_actions = build_unload_actions(&profile, provider_endpoint.as_deref(), provider_pid);
LocalModelLease {
lease_id,
mode: lease_mode_for_profile(&profile),
unload_strategy: profile.unload_strategy.clone(),
profile,
pre_load_gpu_probe,
provider_endpoint,
provider_pid,
idle_ttl_seconds,
unload_actions,
}
}
pub fn build_unload_actions(
profile: &ModelProfile,
provider_endpoint: Option<&str>,
provider_pid: Option<u32>,
) -> Vec<ProviderUnloadAction> {
match profile.unload_strategy {
UnloadStrategy::LlamaCppRouterUnload => {
let endpoint = resolve_provider_endpoint(&profile.unload_strategy, provider_endpoint);
let mut actions = vec![ProviderUnloadAction {
kind: UnloadActionKind::ProviderApi,
label: "llama.cpp router unload".to_string(),
command: None,
http_method: Some("POST".to_string()),
endpoint: Some(endpoint),
body_json: Some(format!(r#"{{"model":"{}"}}"#, profile.model_ref)),
required: true,
}];
if let Some(pid) = provider_pid {
actions.push(process_exit_action(
pid,
"terminate llama.cpp worker if unload is not proven",
));
}
actions
}
UnloadStrategy::OllamaKeepAliveZero => vec![
ProviderUnloadAction {
kind: UnloadActionKind::ProviderApi,
label: "ollama keep_alive zero".to_string(),
command: None,
http_method: Some("POST".to_string()),
endpoint: Some(resolve_provider_endpoint(
&profile.unload_strategy,
provider_endpoint,
)),
body_json: Some(format!(
r#"{{"model":"{}","prompt":"","keep_alive":0}}"#,
profile.model_ref
)),
required: true,
},
ProviderUnloadAction {
kind: UnloadActionKind::ProviderApi,
label: "ollama stop fallback".to_string(),
command: Some(vec![
"ollama".to_string(),
"stop".to_string(),
profile.model_ref.to_string(),
]),
http_method: None,
endpoint: None,
body_json: None,
required: false,
},
],
UnloadStrategy::VllmSleep => {
let endpoint = resolve_provider_endpoint(&profile.unload_strategy, provider_endpoint);
vec![ProviderUnloadAction {
kind: UnloadActionKind::Sleep,
label: "vLLM sleep mode".to_string(),
command: None,
http_method: Some("POST".to_string()),
endpoint: Some(endpoint),
body_json: None,
required: true,
}]
}
UnloadStrategy::ProcessExit => provider_pid
.map(|pid| vec![process_exit_action(pid, "terminate isolated model worker")])
.unwrap_or_else(|| {
vec![ProviderUnloadAction {
kind: UnloadActionKind::ProcessExit,
label: "terminate isolated model worker".to_string(),
command: None,
http_method: None,
endpoint: None,
body_json: None,
required: true,
}]
}),
UnloadStrategy::None => vec![ProviderUnloadAction {
kind: UnloadActionKind::Noop,
label: "no GPU unload required".to_string(),
command: None,
http_method: None,
endpoint: None,
body_json: None,
required: false,
}],
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct UnloadActionResult {
pub label: String,
pub executed: bool,
pub outcome: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PreparedUnloadRequest {
pub label: String,
pub url: String,
pub body: String,
pub fallback_command: Option<Vec<String>>,
}
pub fn rewrite_unload_body_model(body_json: &str, resolved_model_tag: &str) -> String {
if resolved_model_tag.is_empty() {
return body_json.to_string();
}
match serde_json::from_str::<serde_json::Value>(body_json) {
Ok(mut value) => {
if let Some(obj) = value.as_object_mut() {
obj.insert(
"model".to_string(),
serde_json::Value::String(resolved_model_tag.to_string()),
);
}
serde_json::to_string(&value).unwrap_or_else(|_| body_json.to_string())
}
Err(_) => body_json.to_string(),
}
}
pub fn prepare_unload_request(
action: &ProviderUnloadAction,
resolved_model_tag: &str,
) -> Option<PreparedUnloadRequest> {
if action.kind != UnloadActionKind::ProviderApi {
return None;
}
let endpoint = action.endpoint.clone().unwrap_or_default();
let url = normalize_unload_url(&endpoint);
let body = match action.body_json.as_deref() {
Some(template) => rewrite_unload_body_model(template, resolved_model_tag),
None => format!(
r#"{{"model":"{resolved_model_tag}","prompt":"","keep_alive":0}}"#
),
};
Some(PreparedUnloadRequest {
label: action.label.clone(),
url,
body,
fallback_command: action.command.clone(),
})
}
pub fn normalize_unload_url(endpoint: &str) -> String {
let trimmed = endpoint.trim_end_matches('/');
if trimmed.ends_with("/api/generate") {
return trimmed.to_string();
}
format!("{}/api/generate", trimmed)
}
pub fn execute_unload_request(req: &PreparedUnloadRequest) -> UnloadActionResult {
match post_unload_http(&req.url, &req.body) {
Ok(status) => UnloadActionResult {
label: req.label.clone(),
executed: true,
outcome: format!("HTTP {status}"),
},
Err(err) => {
if let Some(cmd) = &req.fallback_command {
let _ = std::process::Command::new(&cmd[0])
.args(&cmd[1..])
.status()
.map_err(|e| {
eprintln!("tsift-local-model: unload fallback command failed: {e}");
});
return UnloadActionResult {
label: req.label.clone(),
executed: true,
outcome: format!("POST failed ({err}); ran fallback `{:?}`", cmd),
};
}
UnloadActionResult {
label: req.label.clone(),
executed: false,
outcome: format!("POST failed ({err}); no fallback"),
}
}
}
}
pub fn unload_model_at(host: &str, model_tag: &str) -> UnloadActionResult {
let req = PreparedUnloadRequest {
label: format!("ollama keep_alive zero for {model_tag}"),
url: normalize_unload_url(host),
body: format!(r#"{{"model":"{model_tag}","prompt":"","keep_alive":0}}"#),
fallback_command: Some(vec![
"ollama".to_string(),
"stop".to_string(),
model_tag.to_string(),
]),
};
execute_unload_request(&req)
}
pub fn execute_unload_actions(
actions: &[ProviderUnloadAction],
resolved_model_tag: &str,
) -> Vec<UnloadActionResult> {
let mut results = Vec::with_capacity(actions.len());
let mut required_succeeded = false;
for action in actions {
if required_succeeded {
results.push(UnloadActionResult {
label: action.label.clone(),
executed: false,
outcome: "skipped: prior required unload succeeded".to_string(),
});
continue;
}
let result = match prepare_unload_request(action, resolved_model_tag) {
Some(req) => execute_unload_request(&req),
None => UnloadActionResult {
label: action.label.clone(),
executed: false,
outcome: "skipped: non-API action".to_string(),
},
};
if result.executed && action.required {
required_succeeded = true;
}
results.push(result);
}
results
}
fn post_unload_http(url: &str, body: &str) -> anyhow::Result<String> {
use std::time::Duration;
let payload = serde_json::from_str::<serde_json::Value>(body)
.unwrap_or(serde_json::Value::Null);
let agent = ureq::Agent::config_builder()
.http_status_as_error(false)
.timeout_global(Some(Duration::from_secs(30)))
.build()
.new_agent();
let mut response = agent
.post(url)
.send_json(payload)
.with_context(|| format!("posting unload to {url}"))?;
let status = response.status();
let text = response
.body_mut()
.read_to_string()
.with_context(|| format!("reading unload response (HTTP {status})"))?;
if !status.is_success() {
bail!("unload HTTP {status}: {}", truncate_str_local(&text, 200));
}
Ok(format!("{status}"))
}
fn truncate_str_local(s: &str, max: usize) -> String {
if s.len() <= max {
s.to_string()
} else {
format!("{}…", &s[..max])
}
}
pub fn resolve_provider_endpoint(strategy: &UnloadStrategy, explicit: Option<&str>) -> String {
if let Some(explicit) = explicit
&& !explicit.trim().is_empty()
{
return explicit.to_string();
}
let (env_var, default): (&str, &str) = match strategy {
UnloadStrategy::LlamaCppRouterUnload => {
(LLAMA_CPP_ENDPOINT_ENV_VAR, DEFAULT_LLAMA_CPP_ENDPOINT)
}
UnloadStrategy::OllamaKeepAliveZero => (OLLAMA_ENDPOINT_ENV_VAR, DEFAULT_OLLAMA_ENDPOINT),
UnloadStrategy::VllmSleep => (VLLM_ENDPOINT_ENV_VAR, DEFAULT_VLLM_ENDPOINT),
UnloadStrategy::ProcessExit | UnloadStrategy::None => return String::new(),
};
if let Ok(value) = std::env::var(env_var)
&& !value.trim().is_empty()
{
return value;
}
default.to_string()
}
pub fn build_lifecycle_report(
profile: ModelProfile,
pre_load_gpu_probe: GpuProbe,
post_unload_gpu_probe: GpuProbe,
provider_endpoint: Option<String>,
provider_pid: Option<u32>,
idle_ttl_seconds: u64,
tolerance_mib: u64,
) -> LocalModelLifecycleReport {
let lease = build_local_model_lease(
profile,
pre_load_gpu_probe.clone(),
provider_endpoint,
provider_pid,
idle_ttl_seconds,
);
let cleanup = evaluate_vram_cleanup(&pre_load_gpu_probe, &post_unload_gpu_probe, tolerance_mib);
let mut notes = vec![match lease.mode {
LeaseMode::Exclusive => {
"large extractor profile requires an exclusive local-model lease".to_string()
}
LeaseMode::Shared => "small model profile can share GPU when the margin fits".to_string(),
LeaseMode::CpuOrHash => "profile does not require GPU VRAM".to_string(),
}];
if !cleanup.cleanup_proven {
notes.push(
"future KG runs should fail if cleanup remains unproven after required unload actions"
.to_string(),
);
}
LocalModelLifecycleReport {
lease,
post_unload_gpu_probe,
cleanup,
notes,
}
}
pub fn evaluate_vram_cleanup(
pre_load_gpu_probe: &GpuProbe,
post_unload_gpu_probe: &GpuProbe,
tolerance_mib: u64,
) -> VramCleanupEvaluation {
let pre_used_mib = pre_load_gpu_probe.used_vram_mib;
let post_used_mib = post_unload_gpu_probe.used_vram_mib;
let allowed_post_used_mib = pre_used_mib.map(|used| used.saturating_add(tolerance_mib));
let used_delta_mib = match (pre_used_mib, post_used_mib) {
(Some(pre), Some(post)) => Some(post as i64 - pre as i64),
_ => None,
};
if !pre_load_gpu_probe.available
|| !post_unload_gpu_probe.available
|| pre_used_mib.is_none()
|| post_used_mib.is_none()
{
return VramCleanupEvaluation {
status: VramCleanupStatus::ProbeUnavailable,
cleanup_proven: false,
pre_used_mib,
post_used_mib,
allowed_post_used_mib,
used_delta_mib,
external_process_delta_mib: 0,
blocking_processes: Vec::new(),
reason: "pre-load or post-unload GPU probe is unavailable".to_string(),
};
}
let pre_used = pre_used_mib.unwrap();
let post_used = post_used_mib.unwrap();
let allowed = allowed_post_used_mib.unwrap();
if post_used <= allowed {
return VramCleanupEvaluation {
status: VramCleanupStatus::Proven,
cleanup_proven: true,
pre_used_mib,
post_used_mib,
allowed_post_used_mib,
used_delta_mib,
external_process_delta_mib: 0,
blocking_processes: Vec::new(),
reason: format!(
"post-unload VRAM {post_used} MiB is within {tolerance_mib} MiB of baseline {pre_used} MiB"
),
};
}
let blocking_processes = post_unload_gpu_probe
.processes
.iter()
.filter(|process| is_tsift_model_process(process))
.cloned()
.collect::<Vec<_>>();
let external_process_delta_mib =
external_process_delta_mib(pre_load_gpu_probe, post_unload_gpu_probe);
if blocking_processes.is_empty()
&& post_used <= allowed.saturating_add(external_process_delta_mib)
{
return VramCleanupEvaluation {
status: VramCleanupStatus::ProvenByExternalAccounting,
cleanup_proven: true,
pre_used_mib,
post_used_mib,
allowed_post_used_mib,
used_delta_mib,
external_process_delta_mib,
blocking_processes,
reason: format!(
"post-unload VRAM increase is accounted for by {external_process_delta_mib} MiB of non-tsift GPU processes"
),
};
}
VramCleanupEvaluation {
status: VramCleanupStatus::NotProven,
cleanup_proven: false,
pre_used_mib,
post_used_mib,
allowed_post_used_mib,
used_delta_mib,
external_process_delta_mib,
blocking_processes,
reason: format!(
"post-unload VRAM {post_used} MiB exceeds allowed {allowed} MiB and cleanup is not externally accounted for"
),
}
}
pub fn rank_profiles_for_role(
profiles: &[ModelProfile],
probe: &GpuProbe,
role: ModelRole,
) -> Vec<ProfileSelection> {
profiles
.iter()
.filter(|profile| profile.supports_role(&role))
.map(|profile| selection_for_profile(profile, probe))
.collect()
}
pub fn format_status_human(report: &LocalModelStatusReport) -> String {
let mut out = String::new();
out.push_str("Local model status\n");
if report.gpu_probe.available {
out.push_str(&format!(
"GPU: {} | VRAM: {} MiB used / {} MiB total ({} MiB free)\n",
report.gpu_probe.gpu_name.as_deref().unwrap_or("unknown"),
format_optional_u64(report.gpu_probe.used_vram_mib),
format_optional_u64(report.gpu_probe.total_vram_mib),
format_optional_u64(report.gpu_probe.free_vram_mib)
));
} else {
out.push_str(&format!(
"GPU: unavailable ({})\n",
report.gpu_probe.error.as_deref().unwrap_or("unknown error")
));
}
out.push_str(&format!(
"Recommended extractor: {}\n",
report
.recommended_extractor
.as_deref()
.unwrap_or("none selectable")
));
out.push_str(&format!(
"Recommended embedding: {}\n",
report
.recommended_embedding
.as_deref()
.unwrap_or("none selectable")
));
out.push_str("\nExtractor profiles:\n");
for selection in &report.extractor_profiles {
out.push_str(&format!(
"- {} [{} MiB est]: {} ({})\n",
selection.profile.id,
selection.profile.estimated_total_mib(),
if selection.selectable {
"selectable"
} else {
"blocked"
},
selection.reason
));
}
out.push_str("\nEmbedding profiles:\n");
for selection in &report.embedding_profiles {
out.push_str(&format!(
"- {} [{} MiB est]: {} ({})\n",
selection.profile.id,
selection.profile.estimated_total_mib(),
if selection.selectable {
"selectable"
} else {
"blocked"
},
selection.reason
));
}
out
}
pub fn format_lifecycle_human(report: &LocalModelLifecycleReport) -> String {
let mut out = String::new();
out.push_str("Local model lifecycle\n");
out.push_str(&format!(
"Profile: {} ({})\n",
report.lease.profile.id, report.lease.profile.label
));
out.push_str(&format!(
"Lease: {} | mode: {:?} | idle TTL: {}s\n",
report.lease.lease_id, report.lease.mode, report.lease.idle_ttl_seconds
));
out.push_str(&format!(
"Pre-load VRAM: {} MiB used\n",
format_optional_u64(report.lease.pre_load_gpu_probe.used_vram_mib)
));
out.push_str(&format!(
"Post-unload VRAM: {} MiB used\n",
format_optional_u64(report.post_unload_gpu_probe.used_vram_mib)
));
out.push_str(&format!(
"Cleanup: {:?} ({})\n",
report.cleanup.status, report.cleanup.reason
));
out.push_str("\nRequired unload actions:\n");
for action in &report.lease.unload_actions {
out.push_str(&format!(
"- {}: {}{}\n",
if action.required {
"required"
} else {
"fallback"
},
action.label,
format_action_detail(action)
));
}
if !report.cleanup.blocking_processes.is_empty() {
out.push_str("\nBlocking GPU processes:\n");
for process in &report.cleanup.blocking_processes {
out.push_str(&format!(
"- pid={} name={} used={} MiB\n",
process
.pid
.map(|pid| pid.to_string())
.unwrap_or_else(|| "unknown".to_string()),
process.process_name,
format_optional_u64(process.used_memory_mib)
));
}
}
out
}
fn format_action_detail(action: &ProviderUnloadAction) -> String {
if let Some(command) = &action.command {
return format!(" | command: {}", command.join(" "));
}
if let Some(endpoint) = &action.endpoint {
return format!(
" | {} {}{}",
action.http_method.as_deref().unwrap_or("POST"),
endpoint,
action
.body_json
.as_ref()
.map(|body| format!(" body={body}"))
.unwrap_or_default()
);
}
String::new()
}
fn selection_for_profile(profile: &ModelProfile, probe: &GpuProbe) -> ProfileSelection {
if profile.concurrency == ConcurrencyClass::CpuOrHash {
return ProfileSelection {
profile: profile.clone(),
selectable: true,
reason: "does not require GPU VRAM".to_string(),
};
}
let Some(free_vram_mib) = probe.free_vram_mib else {
return ProfileSelection {
profile: profile.clone(),
selectable: false,
reason: "free VRAM unknown".to_string(),
};
};
let required = profile.estimated_total_mib();
if required <= free_vram_mib {
ProfileSelection {
profile: profile.clone(),
selectable: true,
reason: format!("estimated {required} MiB fits in {free_vram_mib} MiB free"),
}
} else {
ProfileSelection {
profile: profile.clone(),
selectable: false,
reason: format!("estimated {required} MiB exceeds {free_vram_mib} MiB free"),
}
}
}
fn process_exit_action(pid: u32, label: &str) -> ProviderUnloadAction {
ProviderUnloadAction {
kind: UnloadActionKind::ProcessExit,
label: label.to_string(),
command: Some(vec![
"kill".to_string(),
"-TERM".to_string(),
pid.to_string(),
]),
http_method: None,
endpoint: None,
body_json: None,
required: false,
}
}
fn external_process_delta_mib(
pre_load_gpu_probe: &GpuProbe,
post_unload_gpu_probe: &GpuProbe,
) -> u64 {
post_unload_gpu_probe
.processes
.iter()
.filter(|process| !is_tsift_model_process(process))
.map(|process| {
let before = matching_pre_process(pre_load_gpu_probe, process)
.and_then(|pre| pre.used_memory_mib)
.unwrap_or(0);
process.used_memory_mib.unwrap_or(0).saturating_sub(before)
})
.sum()
}
fn matching_pre_process<'a>(
pre_load_gpu_probe: &'a GpuProbe,
post_process: &GpuProcess,
) -> Option<&'a GpuProcess> {
if let Some(pid) = post_process.pid
&& let Some(process) = pre_load_gpu_probe
.processes
.iter()
.find(|candidate| candidate.pid == Some(pid))
{
return Some(process);
}
pre_load_gpu_probe
.processes
.iter()
.find(|candidate| candidate.process_name == post_process.process_name)
}
fn is_tsift_model_process(process: &GpuProcess) -> bool {
let name = process.process_name.to_ascii_lowercase();
name.contains("tsift")
|| name.contains("llama")
|| name.contains("ollama")
|| name.contains("vllm")
|| name.contains("ggml")
}
pub fn current_unix_seconds() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or(0)
}
fn parse_gpu_query(line: &str) -> Result<GpuProbe> {
let parts = line.split(',').map(str::trim).collect::<Vec<_>>();
if parts.len() != 5 {
anyhow::bail!("unexpected nvidia-smi gpu query row: {line}");
}
Ok(GpuProbe {
timestamp_unix_seconds: Some(current_unix_seconds()),
available: true,
gpu_name: Some(parts[0].to_string()),
driver_version: Some(parts[1].to_string()),
total_vram_mib: Some(parse_optional_u64(parts[2]).context("parse total VRAM")?),
used_vram_mib: Some(parse_optional_u64(parts[3]).context("parse used VRAM")?),
free_vram_mib: Some(parse_optional_u64(parts[4]).context("parse free VRAM")?),
processes: Vec::new(),
error: None,
})
}
fn query_nvidia_compute_processes() -> Vec<GpuProcess> {
let Ok(output) = Command::new("nvidia-smi")
.args([
"--query-compute-apps=pid,process_name,used_memory",
"--format=csv,noheader,nounits",
])
.output()
else {
return Vec::new();
};
if !output.status.success() {
return Vec::new();
}
String::from_utf8_lossy(&output.stdout)
.lines()
.filter_map(parse_process_query)
.collect()
}
fn parse_process_query(line: &str) -> Option<GpuProcess> {
let parts = line.split(',').map(str::trim).collect::<Vec<_>>();
if parts.len() != 3 || parts.iter().all(|part| part.is_empty()) {
return None;
}
Some(GpuProcess {
pid: parts[0].parse::<u32>().ok(),
process_name: parts[1].to_string(),
used_memory_mib: parse_optional_u64(parts[2]).ok(),
})
}
fn parse_optional_u64(input: &str) -> Result<u64> {
let cleaned = input.trim().trim_end_matches("MiB").trim();
cleaned
.parse::<u64>()
.with_context(|| format!("parse integer from {input:?}"))
}
fn format_optional_u64(value: Option<u64>) -> String {
value
.map(|value| value.to_string())
.unwrap_or_else(|| "unknown".to_string())
}
pub const LEASE_REGISTRY_VERSION: u32 = 1;
pub const DEFAULT_LEASE_TTL_SECONDS: u64 = 0;
pub const LEASE_FILE_ENV_VAR: &str = "TSIFT_LEASE_FILE";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GpuLeaseRecord {
pub profile_id: String,
pub holder_pid: u32,
pub holder_command: String,
pub acquired_at_unix_seconds: u64,
pub lease_mode: LeaseMode,
pub vram_baseline_mib: u64,
pub idle_ttl_seconds: u64,
pub notes: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GpuLeaseRegistry {
pub version: u32,
pub leases: BTreeMap<String, Vec<GpuLeaseRecord>>,
}
impl Default for GpuLeaseRegistry {
fn default() -> Self {
Self {
version: LEASE_REGISTRY_VERSION,
leases: BTreeMap::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum GpuLeaseAcquisitionStatus {
Acquired,
Refreshed,
ReclaimedStale,
CpuOrHashBypass,
Conflict,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuLeaseConflict {
pub profile_id: String,
pub holder_pid: u32,
pub holder_command: String,
pub acquired_at_unix_seconds: u64,
pub lease_mode: LeaseMode,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuLeaseAcquisition {
pub profile_id: String,
pub holder_pid: u32,
pub status: GpuLeaseAcquisitionStatus,
pub record: Option<GpuLeaseRecord>,
pub conflict: Option<GpuLeaseConflict>,
pub reclaimed: Vec<GpuLeaseRecord>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum GpuLeaseReleaseOutcome {
Released,
NotHeld,
ProfileAbsent,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuLeaseRelease {
pub profile_id: String,
pub holder_pid: u32,
pub outcome: GpuLeaseReleaseOutcome,
pub remaining_holders: u32,
}
pub fn resolve_lease_file(override_path: Option<&Path>) -> PathBuf {
if let Some(path) = override_path {
return path.to_path_buf();
}
if let Ok(path) = std::env::var(LEASE_FILE_ENV_VAR) {
return PathBuf::from(path);
}
if let Ok(state_dir) = std::env::var("XDG_STATE_HOME")
&& !state_dir.is_empty()
{
return PathBuf::from(state_dir)
.join("tsift")
.join("gpu-lease.json");
}
if let Ok(home) = std::env::var("HOME")
&& !home.is_empty()
{
return PathBuf::from(home).join(".tsift").join("gpu-lease.json");
}
PathBuf::from("./.tsift/gpu-lease.json")
}
pub fn is_pid_alive(pid: u32) -> bool {
if pid == 0 {
return false;
}
if pid == std::process::id() {
return true;
}
match Command::new("kill").arg("-0").arg(pid.to_string()).output() {
Ok(output) => output.status.success(),
Err(_) => false,
}
}
pub fn registry_lock_path(path: &Path) -> PathBuf {
let mut name = path.as_os_str().to_os_string();
name.push(".lock");
PathBuf::from(name)
}
fn with_registry_lock<T>(path: &Path, op: impl FnOnce() -> Result<T>) -> Result<T> {
use fs4::fs_std::FileExt;
let lock_path = registry_lock_path(path);
if let Some(parent) = lock_path.parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent).context("create lease registry lock parent")?;
}
let lock_file = fs::OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(false)
.open(&lock_path)
.with_context(|| format!("open lease registry lock {}", lock_path.display()))?;
lock_file
.lock_exclusive()
.context("acquire exclusive lease registry lock")?;
let result = op();
let _ = FileExt::unlock(&lock_file);
result
}
pub fn read_lease_registry(path: &Path) -> Result<GpuLeaseRegistry> {
match fs::read_to_string(path) {
Ok(contents) => {
if contents.trim().is_empty() {
return Ok(GpuLeaseRegistry::default());
}
serde_json::from_str(&contents).context("parse gpu lease registry")
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
Ok(GpuLeaseRegistry::default())
}
Err(error) => Err(error).context("read gpu lease registry"),
}
}
pub fn write_lease_registry(path: &Path, registry: &GpuLeaseRegistry) -> Result<()> {
if let Some(parent) = path.parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent).context("create lease registry parent")?;
}
let payload = serde_json::to_string_pretty(registry).context("serialize lease registry")?;
let temp_path = path.with_extension(format!(
"json.tmp.{}.{}",
std::process::id(),
current_unix_seconds()
));
let mut handle = fs::File::create(&temp_path).context("create lease registry temp file")?;
handle
.write_all(payload.as_bytes())
.context("write lease registry temp file")?;
handle.sync_all().context("sync lease registry temp file")?;
drop(handle);
fs::rename(&temp_path, path).context("rename lease registry into place")?;
Ok(())
}
pub fn prune_stale_leases(
registry: &mut GpuLeaseRegistry,
now: u64,
is_alive: impl Fn(u32) -> bool,
) -> Vec<GpuLeaseRecord> {
let mut pruned = Vec::new();
let mut empty_keys = Vec::new();
for (profile_id, holders) in registry.leases.iter_mut() {
let mut kept = Vec::with_capacity(holders.len());
for record in holders.drain(..) {
let pid_dead = !is_alive(record.holder_pid);
let ttl_expired = record.idle_ttl_seconds > 0
&& now.saturating_sub(record.acquired_at_unix_seconds) > record.idle_ttl_seconds;
if pid_dead || ttl_expired {
pruned.push(record);
} else {
kept.push(record);
}
}
if kept.is_empty() {
empty_keys.push(profile_id.clone());
}
*holders = kept;
}
for key in empty_keys {
registry.leases.remove(&key);
}
pruned
}
#[allow(clippy::too_many_arguments)]
pub fn apply_acquire(
registry: &mut GpuLeaseRegistry,
profile: &ModelProfile,
holder_pid: u32,
holder_command: &str,
vram_baseline_mib: u64,
idle_ttl_seconds: u64,
now: u64,
is_alive: impl Fn(u32) -> bool,
) -> GpuLeaseAcquisition {
if profile.concurrency == ConcurrencyClass::CpuOrHash {
return GpuLeaseAcquisition {
profile_id: profile.id.to_string(),
holder_pid,
status: GpuLeaseAcquisitionStatus::CpuOrHashBypass,
record: None,
conflict: None,
reclaimed: Vec::new(),
};
}
let reclaimed = prune_stale_leases(registry, now, &is_alive);
let mode = lease_mode_for_profile(profile);
let entry = registry.leases.entry(profile.id.to_string()).or_default();
let already_held = entry
.iter()
.position(|record| record.holder_pid == holder_pid);
let record = GpuLeaseRecord {
profile_id: profile.id.to_string(),
holder_pid,
holder_command: holder_command.to_string(),
acquired_at_unix_seconds: now,
lease_mode: mode.clone(),
vram_baseline_mib,
idle_ttl_seconds,
notes: Vec::new(),
};
let status = if let Some(index) = already_held {
entry[index] = record.clone();
GpuLeaseAcquisitionStatus::Refreshed
} else {
match mode {
LeaseMode::Exclusive => {
if let Some(blocker) = entry.first() {
return GpuLeaseAcquisition {
profile_id: profile.id.to_string(),
holder_pid,
status: GpuLeaseAcquisitionStatus::Conflict,
record: None,
conflict: Some(GpuLeaseConflict {
profile_id: profile.id.to_string(),
holder_pid: blocker.holder_pid,
holder_command: blocker.holder_command.clone(),
acquired_at_unix_seconds: blocker.acquired_at_unix_seconds,
lease_mode: blocker.lease_mode.clone(),
}),
reclaimed,
};
}
entry.push(record.clone());
if reclaimed
.iter()
.any(|pruned| pruned.profile_id == profile.id)
{
GpuLeaseAcquisitionStatus::ReclaimedStale
} else {
GpuLeaseAcquisitionStatus::Acquired
}
}
LeaseMode::Shared => {
entry.push(record.clone());
if reclaimed
.iter()
.any(|pruned| pruned.profile_id == profile.id)
{
GpuLeaseAcquisitionStatus::ReclaimedStale
} else {
GpuLeaseAcquisitionStatus::Acquired
}
}
LeaseMode::CpuOrHash => GpuLeaseAcquisitionStatus::CpuOrHashBypass,
}
};
GpuLeaseAcquisition {
profile_id: profile.id.to_string(),
holder_pid,
status,
record: Some(record),
conflict: None,
reclaimed,
}
}
pub fn apply_release(
registry: &mut GpuLeaseRegistry,
profile_id: &str,
holder_pid: u32,
now: u64,
is_alive: impl Fn(u32) -> bool,
) -> GpuLeaseRelease {
let _ = prune_stale_leases(registry, now, &is_alive);
let Some(holders) = registry.leases.get_mut(profile_id) else {
return GpuLeaseRelease {
profile_id: profile_id.to_string(),
holder_pid,
outcome: GpuLeaseReleaseOutcome::ProfileAbsent,
remaining_holders: 0,
};
};
let before = holders.len();
holders.retain(|record| record.holder_pid != holder_pid);
let removed = before - holders.len();
let remaining = holders.len() as u32;
if holders.is_empty() {
registry.leases.remove(profile_id);
}
let outcome = if removed == 0 {
GpuLeaseReleaseOutcome::NotHeld
} else {
GpuLeaseReleaseOutcome::Released
};
GpuLeaseRelease {
profile_id: profile_id.to_string(),
holder_pid,
outcome,
remaining_holders: remaining,
}
}
pub fn acquire_lease(
profile_id: &str,
holder_pid: u32,
holder_command: &str,
vram_baseline_mib: u64,
idle_ttl_seconds: u64,
now: u64,
path: &Path,
) -> Result<GpuLeaseAcquisition> {
let profile = profile_by_id(profile_id)
.with_context(|| format!("unknown local model profile {profile_id:?}"))?;
with_registry_lock(path, || {
let mut registry = read_lease_registry(path)?;
let acquisition = apply_acquire(
&mut registry,
&profile,
holder_pid,
holder_command,
vram_baseline_mib,
idle_ttl_seconds,
now,
is_pid_alive,
);
if acquisition.status != GpuLeaseAcquisitionStatus::CpuOrHashBypass {
write_lease_registry(path, ®istry)?;
}
Ok(acquisition)
})
}
pub fn release_lease(
profile_id: &str,
holder_pid: u32,
now: u64,
path: &Path,
) -> Result<GpuLeaseRelease> {
with_registry_lock(path, || {
let mut registry = read_lease_registry(path)?;
let release = apply_release(&mut registry, profile_id, holder_pid, now, is_pid_alive);
write_lease_registry(path, ®istry)?;
Ok(release)
})
}
pub fn show_registry(path: &Path, now: u64, include_stale: bool) -> Result<GpuLeaseRegistry> {
let mut registry = read_lease_registry(path)?;
if !include_stale {
prune_stale_leases(&mut registry, now, is_pid_alive);
}
Ok(registry)
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum GpuLeaseRenewOutcome {
Renewed,
NotHeld,
ProfileAbsent,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuLeaseRenew {
pub profile_id: String,
pub holder_pid: u32,
pub outcome: GpuLeaseRenewOutcome,
pub renewed_at_unix_seconds: Option<u64>,
}
pub fn apply_renew(
registry: &mut GpuLeaseRegistry,
profile_id: &str,
holder_pid: u32,
now: u64,
is_alive: impl Fn(u32) -> bool,
) -> GpuLeaseRenew {
let _ = prune_stale_leases(registry, now, &is_alive);
let Some(holders) = registry.leases.get_mut(profile_id) else {
return GpuLeaseRenew {
profile_id: profile_id.to_string(),
holder_pid,
outcome: GpuLeaseRenewOutcome::ProfileAbsent,
renewed_at_unix_seconds: None,
};
};
if let Some(record) = holders
.iter_mut()
.find(|record| record.holder_pid == holder_pid)
{
record.acquired_at_unix_seconds = now;
GpuLeaseRenew {
profile_id: profile_id.to_string(),
holder_pid,
outcome: GpuLeaseRenewOutcome::Renewed,
renewed_at_unix_seconds: Some(now),
}
} else {
GpuLeaseRenew {
profile_id: profile_id.to_string(),
holder_pid,
outcome: GpuLeaseRenewOutcome::NotHeld,
renewed_at_unix_seconds: None,
}
}
}
pub fn renew_lease(
profile_id: &str,
holder_pid: u32,
now: u64,
path: &Path,
) -> Result<GpuLeaseRenew> {
with_registry_lock(path, || {
let mut registry = read_lease_registry(path)?;
let renew = apply_renew(&mut registry, profile_id, holder_pid, now, is_pid_alive);
write_lease_registry(path, ®istry)?;
Ok(renew)
})
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct GpuLeaseReap {
pub reclaimed: Vec<GpuLeaseRecord>,
pub emptied_profiles: Vec<String>,
}
pub fn reap_leases(now: u64, path: &Path) -> Result<GpuLeaseReap> {
with_registry_lock(path, || {
let mut registry = read_lease_registry(path)?;
let before: std::collections::BTreeSet<String> =
registry.leases.keys().cloned().collect();
let reclaimed = prune_stale_leases(&mut registry, now, is_pid_alive);
let emptied_profiles: Vec<String> = before
.into_iter()
.filter(|profile_id| !registry.leases.contains_key(profile_id))
.collect();
write_lease_registry(path, ®istry)?;
Ok(GpuLeaseReap {
reclaimed,
emptied_profiles,
})
})
}
pub fn format_lease_show_human(registry: &GpuLeaseRegistry, now: u64) -> String {
let mut out = String::new();
out.push_str("GPU lease registry\n");
out.push_str(&format!("version: {}\n", registry.version));
if registry.leases.is_empty() {
out.push_str("leases: none\n");
return out;
}
out.push_str(&format!("profiles held: {}\n", registry.leases.len()));
for (profile_id, holders) in ®istry.leases {
out.push_str(&format!("\n{profile_id} ({} holder(s)):\n", holders.len()));
for record in holders {
let age = now.saturating_sub(record.acquired_at_unix_seconds);
out.push_str(&format!(
" pid={} cmd={} mode={:?} acquired={}s ago baseline={} MiB ttl={}s",
record.holder_pid,
record.holder_command,
record.lease_mode,
age,
record.vram_baseline_mib,
record.idle_ttl_seconds
));
if record.notes.is_empty() {
out.push('\n');
} else {
out.push_str(&format!(" notes={}\n", record.notes.join("; ")));
}
}
}
out
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
pub enum ProfilePreference {
Auto,
Pinned(String),
ForceHash,
}
impl ProfilePreference {
pub fn from_cli(value: Option<&str>) -> Self {
match value.map(str::trim) {
None | Some("") => ProfilePreference::Auto,
Some("hash") | Some("tsift-local-hash-v1") => ProfilePreference::ForceHash,
Some(other) => ProfilePreference::Pinned(other.to_string()),
}
}
pub fn describe(&self) -> String {
match self {
ProfilePreference::Auto => "auto".to_string(),
ProfilePreference::Pinned(id) => format!("pinned:{id}"),
ProfilePreference::ForceHash => "force-hash".to_string(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ProfileResolutionSource {
AutoRanked,
Pinned,
PinnedUnselectable,
ForcedHash,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ProfileResolution {
pub preference: ProfilePreference,
pub role: ModelRole,
pub source: ProfileResolutionSource,
pub profile: ModelProfile,
pub selectable: bool,
pub reason: String,
}
pub fn resolve_profile_preference(
preference: &ProfilePreference,
role: ModelRole,
probe: &GpuProbe,
) -> ProfileResolution {
let profiles = default_model_profiles();
let hash_profile = profiles
.iter()
.find(|profile| profile.id == "tsift-local-hash-v1")
.cloned()
.expect("hash fallback profile is always present");
match preference {
ProfilePreference::ForceHash => ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::ForcedHash,
profile: hash_profile,
selectable: true,
reason: "caller forced the CPU/hash fallback".to_string(),
},
ProfilePreference::Auto => {
let ranked = rank_profiles_for_role(&profiles, probe, role);
let pick = ranked
.iter()
.find(|selection| selection.selectable)
.cloned()
.or_else(|| {
ranked.into_iter().next().map(|selection| ProfileSelection {
selectable: false,
..selection
})
});
match pick {
Some(selection) if selection.selectable => ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::AutoRanked,
profile: selection.profile.clone(),
selectable: true,
reason: format!("auto-ranked: {}", selection.reason),
},
Some(selection) => ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::AutoRanked,
profile: hash_profile,
selectable: true,
reason: format!(
"auto-ranked but no GPU profile selectable ({}); using hash fallback",
selection.reason
),
},
None => ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::AutoRanked,
profile: hash_profile,
selectable: true,
reason: "no profile matches the requested role; using hash fallback"
.to_string(),
},
}
}
ProfilePreference::Pinned(id) => match profile_by_id(id) {
Some(profile) if profile.supports_role(&role) => {
let selection = selection_for_profile(&profile, probe);
if selection.selectable {
ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::Pinned,
profile,
selectable: true,
reason: format!("pinned: {}", selection.reason),
}
} else {
ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::PinnedUnselectable,
profile: hash_profile,
selectable: true,
reason: format!(
"pinned {} is not selectable ({}); using hash fallback",
id, selection.reason
),
}
}
}
Some(_) => ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::PinnedUnselectable,
profile: hash_profile,
selectable: true,
reason: format!(
"pinned {id} does not support role {:?}; using hash fallback",
role
),
},
None => ProfileResolution {
preference: preference.clone(),
role,
source: ProfileResolutionSource::PinnedUnselectable,
profile: hash_profile,
selectable: true,
reason: format!("pinned profile id {id:?} is unknown; using hash fallback"),
},
},
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SwapStatus {
Swapped,
SwappedToHash,
UnloadProvenTargetUnselectable,
UnloadNotProven,
NoOpSameProfile,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct LocalModelSwapReport {
pub from_profile_id: String,
pub to_profile_id: String,
pub unload: LocalModelLifecycleReport,
pub target_resolution: ProfileResolution,
pub swap_status: SwapStatus,
pub notes: Vec<String>,
}
#[allow(clippy::too_many_arguments)]
pub fn build_swap_report(
from_profile: ModelProfile,
to_profile: ModelProfile,
pre_load_probe: GpuProbe,
post_unload_probe: GpuProbe,
provider_endpoint: Option<String>,
provider_pid: Option<u32>,
idle_ttl_seconds: u64,
tolerance_mib: u64,
) -> LocalModelSwapReport {
let unload = build_lifecycle_report(
from_profile.clone(),
pre_load_probe,
post_unload_probe.clone(),
provider_endpoint,
provider_pid,
idle_ttl_seconds,
tolerance_mib,
);
let target_role = to_profile
.roles
.first()
.copied()
.unwrap_or(ModelRole::Extract);
let target_resolution = resolve_profile_preference(
&ProfilePreference::Pinned(to_profile.id.to_string()),
target_role,
&post_unload_probe,
);
let swap_status = if from_profile.id == to_profile.id {
SwapStatus::NoOpSameProfile
} else if !unload.cleanup.cleanup_proven {
SwapStatus::UnloadNotProven
} else if to_profile.concurrency == ConcurrencyClass::CpuOrHash {
SwapStatus::SwappedToHash
} else if target_resolution.selectable && target_resolution.profile.id == to_profile.id {
SwapStatus::Swapped
} else {
SwapStatus::UnloadProvenTargetUnselectable
};
let mut notes = vec![
format!("swapping from {} to {}", from_profile.id, to_profile.id),
format!("unload cleanup: {:?}", unload.cleanup.status),
format!("target resolution: {:?}", target_resolution.source),
];
if swap_status == SwapStatus::UnloadNotProven {
notes.push("DO NOT load target — source unload did not prove VRAM cleanup".to_string());
}
if swap_status == SwapStatus::UnloadProvenTargetUnselectable {
notes.push(format!(
"target {} is not selectable on the post-unload probe; consider the hash fallback or a smaller profile",
to_profile.id
));
}
LocalModelSwapReport {
from_profile_id: from_profile.id.to_string(),
to_profile_id: to_profile.id.to_string(),
unload,
target_resolution,
swap_status,
notes,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn rtx_5090_probe() -> GpuProbe {
GpuProbe {
timestamp_unix_seconds: Some(1_781_000_000),
available: true,
gpu_name: Some("NVIDIA GeForce RTX 5090".to_string()),
driver_version: Some("610.43.02".to_string()),
total_vram_mib: Some(32_607),
used_vram_mib: Some(179),
free_vram_mib: Some(32_428),
processes: Vec::new(),
error: None,
}
}
fn probe_with_used_vram(used_vram_mib: u64) -> GpuProbe {
let mut probe = rtx_5090_probe();
probe.used_vram_mib = Some(used_vram_mib);
probe.free_vram_mib = probe
.total_vram_mib
.map(|total| total.saturating_sub(used_vram_mib));
probe
}
#[test]
fn qwen3_32b_is_default_extractor_for_clear_5090() {
let report = build_status_report_with_probe(rtx_5090_probe());
assert_eq!(
report.recommended_extractor.as_deref(),
Some("qwen3-32b-q4")
);
assert!(
report
.extractor_profiles
.iter()
.any(|selection| selection.profile.id == "qwen3.5-35b-a3b-q4"
&& !selection.selectable)
);
}
#[test]
fn hash_fallback_selects_without_gpu_probe() {
let report = build_status_report_with_probe(GpuProbe::unavailable("missing"));
assert_eq!(
report.recommended_embedding.as_deref(),
Some("tsift-local-hash-v1")
);
assert_eq!(report.recommended_extractor, None);
}
#[test]
fn parses_nvidia_smi_gpu_query_row() {
let probe =
parse_gpu_query("NVIDIA GeForce RTX 5090, 610.43.02, 32607, 179, 32428").unwrap();
assert!(probe.timestamp_unix_seconds.is_some());
assert_eq!(probe.gpu_name.as_deref(), Some("NVIDIA GeForce RTX 5090"));
assert_eq!(probe.total_vram_mib, Some(32_607));
assert_eq!(probe.used_vram_mib, Some(179));
assert_eq!(probe.free_vram_mib, Some(32_428));
}
#[test]
fn lifecycle_report_plans_llamacpp_unload_and_proves_cleanup() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let report = build_lifecycle_report(
profile,
probe_with_used_vram(200),
probe_with_used_vram(820),
Some("http://127.0.0.1:8080/models/unload".to_string()),
Some(42),
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(report.lease.mode, LeaseMode::Exclusive);
assert!(report.cleanup.cleanup_proven);
assert_eq!(report.cleanup.status, VramCleanupStatus::Proven);
assert!(report.lease.unload_actions.iter().any(|action| {
action.kind == UnloadActionKind::ProviderApi
&& action.endpoint.as_deref() == Some("http://127.0.0.1:8080/models/unload")
}));
assert!(report.lease.unload_actions.iter().any(|action| {
action.kind == UnloadActionKind::ProcessExit
&& action.command.as_ref().is_some_and(|command| {
command == &vec!["kill".to_string(), "-TERM".to_string(), "42".to_string()]
})
}));
}
#[test]
fn vram_cleanup_fails_when_provider_process_remains_loaded() {
let pre = probe_with_used_vram(200);
let mut post = probe_with_used_vram(4_000);
post.processes.push(GpuProcess {
pid: Some(42),
process_name: "llama-server".to_string(),
used_memory_mib: Some(3_000),
});
let cleanup = evaluate_vram_cleanup(&pre, &post, DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB);
assert!(!cleanup.cleanup_proven);
assert_eq!(cleanup.status, VramCleanupStatus::NotProven);
assert_eq!(cleanup.blocking_processes.len(), 1);
}
#[test]
fn vram_cleanup_accepts_external_process_accounting() {
let pre = probe_with_used_vram(200);
let mut post = probe_with_used_vram(2_000);
post.processes.push(GpuProcess {
pid: Some(77),
process_name: "python-training-job".to_string(),
used_memory_mib: Some(1_600),
});
let cleanup = evaluate_vram_cleanup(&pre, &post, DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB);
assert!(cleanup.cleanup_proven);
assert_eq!(
cleanup.status,
VramCleanupStatus::ProvenByExternalAccounting
);
assert_eq!(cleanup.external_process_delta_mib, 1_600);
}
#[test]
fn interrupted_run_cleanup_fails_with_orphaned_provider_process() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let pre = probe_with_used_vram(200);
let mut post = probe_with_used_vram(8_000);
post.processes.push(GpuProcess {
pid: Some(1234),
process_name: "ollama runner".to_string(),
used_memory_mib: Some(7_000),
});
let report = build_lifecycle_report(
profile,
pre,
post,
None,
Some(1234),
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert!(!report.cleanup.cleanup_proven);
assert_eq!(report.cleanup.status, VramCleanupStatus::NotProven);
assert_eq!(report.cleanup.blocking_processes.len(), 1);
assert!(report.lease.unload_actions.iter().any(|action| {
action.kind == UnloadActionKind::ProcessExit
&& action.command.as_ref().is_some_and(|command| {
command == &vec!["kill".to_string(), "-TERM".to_string(), "1234".to_string()]
})
}));
}
fn all_alive(_pid: u32) -> bool {
true
}
fn alive_set(alive: &[u32]) -> impl Fn(u32) -> bool + '_ {
move |pid| alive.contains(&pid)
}
#[test]
fn resolve_lease_file_prefers_explicit_override() {
let path = resolve_lease_file(Some(Path::new("/custom/lease.json")));
assert_eq!(path, PathBuf::from("/custom/lease.json"));
}
#[test]
fn resolve_lease_file_returns_env_value_when_set() {
unsafe {
std::env::set_var(LEASE_FILE_ENV_VAR, "/env/lease.json");
}
let path = resolve_lease_file(None);
unsafe {
std::env::remove_var(LEASE_FILE_ENV_VAR);
}
assert_eq!(path, PathBuf::from("/env/lease.json"));
}
#[test]
fn lease_registry_round_trips_through_json() {
let mut registry = GpuLeaseRegistry::default();
registry.leases.insert(
"qwen3-32b-q4".to_string(),
vec![GpuLeaseRecord {
profile_id: "qwen3-32b-q4".to_string(),
holder_pid: 4242,
holder_command: "tsift".to_string(),
acquired_at_unix_seconds: 100,
lease_mode: LeaseMode::Exclusive,
vram_baseline_mib: 200,
idle_ttl_seconds: 0,
notes: vec!["baseline".to_string()],
}],
);
let payload = serde_json::to_string(®istry).unwrap();
let back: GpuLeaseRegistry = serde_json::from_str(&payload).unwrap();
assert_eq!(registry, back);
assert_eq!(back.version, LEASE_REGISTRY_VERSION);
}
#[test]
fn acquire_exclusive_profile_succeeds_when_free() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
let acquisition = apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
assert_eq!(acquisition.status, GpuLeaseAcquisitionStatus::Acquired);
assert!(acquisition.conflict.is_none());
assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 100);
}
#[test]
fn acquire_exclusive_profile_conflicts_with_live_holder() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
let second = apply_acquire(
&mut registry,
&profile,
200,
"corky",
250,
0,
1_050,
alive_set(&[100, 200]),
);
assert_eq!(second.status, GpuLeaseAcquisitionStatus::Conflict);
let conflict = second.conflict.unwrap();
assert_eq!(conflict.holder_pid, 100);
assert_eq!(conflict.holder_command, "tsift");
assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 100);
}
#[test]
fn acquire_exclusive_profile_reclaims_when_holder_pid_dead() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
let reclaimed = apply_acquire(
&mut registry,
&profile,
200,
"corky",
250,
0,
1_050,
alive_set(&[200]),
);
assert_eq!(reclaimed.status, GpuLeaseAcquisitionStatus::ReclaimedStale);
assert_eq!(reclaimed.reclaimed.len(), 1);
assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 200);
}
#[test]
fn acquire_shared_profile_allows_multiple_live_holders() {
let profile = profile_by_id("qwen3-embedding-0.6b").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
let second = apply_acquire(
&mut registry,
&profile,
200,
"headroom",
250,
0,
1_050,
alive_set(&[100, 200]),
);
assert_eq!(second.status, GpuLeaseAcquisitionStatus::Acquired);
assert_eq!(registry.leases["qwen3-embedding-0.6b"].len(), 2);
}
#[test]
fn acquire_refreshes_when_same_holder_requests_again() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
let again = apply_acquire(
&mut registry,
&profile,
100,
"tsift",
180,
0,
1_500,
all_alive,
);
assert_eq!(again.status, GpuLeaseAcquisitionStatus::Refreshed);
assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
assert_eq!(
registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
1_500
);
assert_eq!(registry.leases["qwen3-32b-q4"][0].vram_baseline_mib, 180);
}
#[test]
fn acquire_cpu_or_hash_profile_bypasses_registry() {
let profile = profile_by_id("tsift-local-hash-v1").unwrap();
let mut registry = GpuLeaseRegistry::default();
let bypass = apply_acquire(
&mut registry,
&profile,
100,
"tsift",
0,
0,
1_000,
all_alive,
);
assert_eq!(bypass.status, GpuLeaseAcquisitionStatus::CpuOrHashBypass);
assert!(registry.leases.is_empty());
}
#[test]
fn idle_ttl_expires_even_when_pid_still_alive() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
60,
1_000,
all_alive,
);
let reclaimed = apply_acquire(
&mut registry,
&profile,
200,
"corky",
250,
0,
1_120,
all_alive,
);
assert_eq!(reclaimed.status, GpuLeaseAcquisitionStatus::ReclaimedStale);
assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 200);
}
#[test]
fn release_removes_holder_and_drops_empty_profile() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
let release = apply_release(&mut registry, "qwen3-32b-q4", 100, 1_050, all_alive);
assert_eq!(release.outcome, GpuLeaseReleaseOutcome::Released);
assert_eq!(release.remaining_holders, 0);
assert!(registry.leases.is_empty());
}
#[test]
fn release_by_non_holder_reports_not_held() {
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
apply_acquire(
&mut registry,
&profile,
100,
"tsift",
200,
0,
1_000,
all_alive,
);
let release = apply_release(&mut registry, "qwen3-32b-q4", 999, 1_050, all_alive);
assert_eq!(release.outcome, GpuLeaseReleaseOutcome::NotHeld);
assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
}
#[test]
fn acquire_and_release_round_trip_through_file() {
let dir = tempfile_dir();
let path = dir.join("gpu-lease.json");
let profile = profile_by_id("qwen3-32b-q4").unwrap();
let mut registry = GpuLeaseRegistry::default();
let acquisition = apply_acquire(
&mut registry,
&profile,
4242,
"tsift",
220,
0,
1_000,
all_alive,
);
assert_eq!(acquisition.status, GpuLeaseAcquisitionStatus::Acquired);
write_lease_registry(&path, ®istry).unwrap();
let read_back = read_lease_registry(&path).unwrap();
assert_eq!(read_back, registry);
assert_eq!(read_back.leases["qwen3-32b-q4"][0].holder_pid, 4242);
let release = apply_release(&mut registry, "qwen3-32b-q4", 4242, 1_050, all_alive);
assert_eq!(release.outcome, GpuLeaseReleaseOutcome::Released);
write_lease_registry(&path, ®istry).unwrap();
let after = read_lease_registry(&path).unwrap();
assert!(after.leases.is_empty());
}
#[test]
fn read_lease_registry_returns_default_for_missing_file() {
let path = Path::new("/definitely/not/a/real/path/lease.json");
let registry = read_lease_registry(path).unwrap();
assert_eq!(registry, GpuLeaseRegistry::default());
}
#[test]
fn registry_lock_path_appends_lock_suffix() {
assert_eq!(
registry_lock_path(Path::new("/tmp/x/gpu-lease.json")),
PathBuf::from("/tmp/x/gpu-lease.json.lock")
);
}
#[test]
fn acquire_lease_creates_sidecar_lock_file() {
let dir = tempfile_dir();
let path = dir.join("gpu-lease.json");
acquire_lease("qwen3-32b-q4", std::process::id(), "tsift", 0, 0, 1_000, &path).unwrap();
assert!(
registry_lock_path(&path).exists(),
"sidecar lock file should exist after a locked acquire"
);
}
#[test]
fn apply_renew_slides_heartbeat_so_ttl_holder_survives() {
let mut registry = GpuLeaseRegistry::default();
let profile = profile_by_id("qwen3-32b-q4").unwrap();
apply_acquire(&mut registry, &profile, 100, "tsift", 200, 100, 1_000, all_alive);
let renew = apply_renew(&mut registry, "qwen3-32b-q4", 100, 1_050, all_alive);
assert_eq!(renew.outcome, GpuLeaseRenewOutcome::Renewed);
assert_eq!(renew.renewed_at_unix_seconds, Some(1_050));
assert_eq!(
registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
1_050
);
let pruned = prune_stale_leases(&mut registry, 1_120, all_alive);
assert!(pruned.is_empty());
assert!(registry.leases.contains_key("qwen3-32b-q4"));
}
#[test]
fn apply_renew_reports_profile_absent_for_unheld_profile() {
let mut registry = GpuLeaseRegistry::default();
let renew = apply_renew(&mut registry, "qwen3-32b-q4", 100, 1_000, all_alive);
assert_eq!(renew.outcome, GpuLeaseRenewOutcome::ProfileAbsent);
assert!(renew.renewed_at_unix_seconds.is_none());
}
#[test]
fn reap_leases_reclaims_dead_pid_and_reports_emptied_profile() {
let dir = tempfile_dir();
let path = dir.join("gpu-lease.json");
let mut registry = GpuLeaseRegistry::default();
registry.leases.insert(
"qwen3-32b-q4".to_string(),
vec![GpuLeaseRecord {
profile_id: "qwen3-32b-q4".to_string(),
holder_pid: 4_000_000_000,
holder_command: "crashed-session".to_string(),
acquired_at_unix_seconds: 1_000,
lease_mode: LeaseMode::Exclusive,
vram_baseline_mib: 200,
idle_ttl_seconds: 0,
notes: Vec::new(),
}],
);
write_lease_registry(&path, ®istry).unwrap();
let reap = reap_leases(2_000, &path).unwrap();
assert_eq!(reap.reclaimed.len(), 1);
assert_eq!(reap.emptied_profiles, vec!["qwen3-32b-q4".to_string()]);
let after = read_lease_registry(&path).unwrap();
assert!(after.leases.is_empty());
}
#[test]
fn renew_lease_round_trips_through_file() {
let dir = tempfile_dir();
let path = dir.join("gpu-lease.json");
let pid = std::process::id();
acquire_lease("qwen3-32b-q4", pid, "tsift", 0, 0, 1_000, &path).unwrap();
let renew = renew_lease("qwen3-32b-q4", pid, 5_000, &path).unwrap();
assert_eq!(renew.outcome, GpuLeaseRenewOutcome::Renewed);
let registry = read_lease_registry(&path).unwrap();
assert_eq!(
registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
5_000
);
}
#[test]
fn prune_stale_leaves_healthy_entries_alone() {
let mut registry = GpuLeaseRegistry::default();
registry.leases.insert(
"qwen3-32b-q4".to_string(),
vec![GpuLeaseRecord {
profile_id: "qwen3-32b-q4".to_string(),
holder_pid: 100,
holder_command: "tsift".to_string(),
acquired_at_unix_seconds: 1_000,
lease_mode: LeaseMode::Exclusive,
vram_baseline_mib: 200,
idle_ttl_seconds: 0,
notes: Vec::new(),
}],
);
let pruned = prune_stale_leases(&mut registry, 1_010, alive_set(&[100]));
assert!(pruned.is_empty());
assert!(registry.leases.contains_key("qwen3-32b-q4"));
}
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
ENV_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn tempfile_dir() -> PathBuf {
use std::sync::atomic::{AtomicU64, Ordering};
static COUNTER: AtomicU64 = AtomicU64::new(0);
let dir = std::env::temp_dir().join(format!(
"tsift-lease-test-{}-{}-{}",
std::process::id(),
current_unix_seconds(),
COUNTER.fetch_add(1, Ordering::Relaxed)
));
std::fs::create_dir_all(&dir).unwrap();
dir
}
#[test]
fn profile_preference_parses_cli_value() {
assert_eq!(ProfilePreference::from_cli(None), ProfilePreference::Auto);
assert_eq!(
ProfilePreference::from_cli(Some("")),
ProfilePreference::Auto
);
assert_eq!(
ProfilePreference::from_cli(Some("hash")),
ProfilePreference::ForceHash
);
assert_eq!(
ProfilePreference::from_cli(Some("tsift-local-hash-v1")),
ProfilePreference::ForceHash
);
assert_eq!(
ProfilePreference::from_cli(Some("qwen3-32b-q4")),
ProfilePreference::Pinned("qwen3-32b-q4".to_string())
);
}
#[test]
fn resolve_auto_picks_recommended_gpu_profile_on_clear_5090() {
let probe = rtx_5090_probe();
let resolution =
resolve_profile_preference(&ProfilePreference::Auto, ModelRole::Extract, &probe);
assert_eq!(resolution.source, ProfileResolutionSource::AutoRanked);
assert!(resolution.selectable);
assert_eq!(resolution.profile.id, "qwen3-32b-q4");
}
#[test]
fn resolve_auto_falls_back_to_hash_when_gpu_unavailable() {
let probe = GpuProbe::unavailable("missing");
let resolution =
resolve_profile_preference(&ProfilePreference::Auto, ModelRole::Extract, &probe);
assert_eq!(resolution.source, ProfileResolutionSource::AutoRanked);
assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
assert!(resolution.selectable);
assert!(resolution.reason.contains("no GPU profile selectable"));
}
#[test]
fn resolve_pinned_selectable_profile_is_used_as_is() {
let probe = rtx_5090_probe();
let resolution = resolve_profile_preference(
&ProfilePreference::Pinned("qwen3-embedding-0.6b".to_string()),
ModelRole::Embed,
&probe,
);
assert_eq!(resolution.source, ProfileResolutionSource::Pinned);
assert_eq!(resolution.profile.id, "qwen3-embedding-0.6b");
assert!(resolution.selectable);
}
#[test]
fn resolve_pinned_profile_with_wrong_role_falls_back_to_hash() {
let probe = rtx_5090_probe();
let resolution = resolve_profile_preference(
&ProfilePreference::Pinned("qwen3-embedding-0.6b".to_string()),
ModelRole::Extract,
&probe,
);
assert_eq!(
resolution.source,
ProfileResolutionSource::PinnedUnselectable
);
assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
assert!(resolution.reason.contains("does not support role"));
}
#[test]
fn resolve_pinned_profile_that_does_not_fit_vram_falls_back_to_hash() {
let probe = probe_with_used_vram(30_000);
let resolution = resolve_profile_preference(
&ProfilePreference::Pinned("qwen3-32b-q4".to_string()),
ModelRole::Extract,
&probe,
);
assert_eq!(
resolution.source,
ProfileResolutionSource::PinnedUnselectable
);
assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
assert!(resolution.selectable);
assert!(resolution.reason.contains("not selectable"));
}
#[test]
fn resolve_force_hash_always_uses_hash_profile() {
let probe = rtx_5090_probe();
let resolution =
resolve_profile_preference(&ProfilePreference::ForceHash, ModelRole::Extract, &probe);
assert_eq!(resolution.source, ProfileResolutionSource::ForcedHash);
assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
assert!(resolution.selectable);
assert!(resolution.reason.contains("forced"));
}
#[test]
fn resolve_pinned_unknown_profile_id_falls_back_to_hash() {
let probe = rtx_5090_probe();
let resolution = resolve_profile_preference(
&ProfilePreference::Pinned("not-a-real-profile".to_string()),
ModelRole::Embed,
&probe,
);
assert_eq!(
resolution.source,
ProfileResolutionSource::PinnedUnselectable
);
assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
assert!(resolution.reason.contains("unknown"));
}
#[test]
fn resolve_endpoint_returns_explicit_override_for_any_strategy() {
for strategy in [
UnloadStrategy::LlamaCppRouterUnload,
UnloadStrategy::OllamaKeepAliveZero,
UnloadStrategy::VllmSleep,
UnloadStrategy::ProcessExit,
UnloadStrategy::None,
] {
let resolved = resolve_provider_endpoint(&strategy, Some("http://custom:9999/path"));
assert_eq!(
resolved, "http://custom:9999/path",
"explicit override should win for {strategy:?}"
);
}
}
#[test]
fn resolve_endpoint_uses_compile_time_default_when_no_env_no_explicit() {
let _env = env_lock();
unsafe {
std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
std::env::remove_var(OLLAMA_ENDPOINT_ENV_VAR);
std::env::remove_var(VLLM_ENDPOINT_ENV_VAR);
}
assert_eq!(
resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None),
DEFAULT_LLAMA_CPP_ENDPOINT
);
assert_eq!(
resolve_provider_endpoint(&UnloadStrategy::OllamaKeepAliveZero, None),
DEFAULT_OLLAMA_ENDPOINT
);
assert_eq!(
resolve_provider_endpoint(&UnloadStrategy::VllmSleep, None),
DEFAULT_VLLM_ENDPOINT
);
assert_eq!(
resolve_provider_endpoint(&UnloadStrategy::ProcessExit, None),
""
);
assert_eq!(resolve_provider_endpoint(&UnloadStrategy::None, None), "");
}
#[test]
fn resolve_endpoint_env_var_overrides_default_for_llama_cpp() {
let _env = env_lock();
unsafe {
std::env::set_var(
LLAMA_CPP_ENDPOINT_ENV_VAR,
"http://127.0.0.1:8081/models/unload",
);
}
let resolved = resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None);
unsafe {
std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
}
assert_eq!(resolved, "http://127.0.0.1:8081/models/unload");
}
#[test]
fn resolve_endpoint_blank_env_var_falls_back_to_default() {
let _env = env_lock();
unsafe {
std::env::set_var(LLAMA_CPP_ENDPOINT_ENV_VAR, " ");
}
let resolved = resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None);
unsafe {
std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
}
assert_eq!(resolved, DEFAULT_LLAMA_CPP_ENDPOINT);
}
#[test]
fn build_unload_actions_picks_up_env_var_for_llama_cpp_endpoint() {
let _env = env_lock();
let profile = profile_by_id("qwen3-32b-q4").unwrap();
unsafe {
std::env::set_var(
LLAMA_CPP_ENDPOINT_ENV_VAR,
"http://127.0.0.1:8081/models/unload",
);
}
let actions = build_unload_actions(&profile, None, Some(42));
unsafe {
std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
}
let unload_action = actions
.iter()
.find(|action| action.kind == UnloadActionKind::ProviderApi)
.expect("provider api action present");
assert_eq!(
unload_action.endpoint.as_deref(),
Some("http://127.0.0.1:8081/models/unload")
);
}
fn probe_pair(pre_used: u64, post_used: u64) -> (GpuProbe, GpuProbe) {
(
probe_with_used_vram(pre_used),
probe_with_used_vram(post_used),
)
}
#[test]
fn swap_to_same_profile_is_noop() {
let from = profile_by_id("qwen3-32b-q4").unwrap();
let to = profile_by_id("qwen3-32b-q4").unwrap();
let (pre, post) = probe_pair(200, 200);
let report = build_swap_report(
from,
to,
pre,
post,
None,
None,
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(report.swap_status, SwapStatus::NoOpSameProfile);
}
#[test]
fn swap_from_big_to_small_embedding_when_cleanup_proven_is_swapped() {
let from = profile_by_id("qwen3-32b-q4").unwrap();
let to = profile_by_id("qwen3-embedding-0.6b").unwrap();
let (pre, post) = probe_pair(28_000, 200);
let report = build_swap_report(
from,
to,
pre,
post,
None,
Some(42),
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(report.swap_status, SwapStatus::Swapped);
assert!(report.unload.cleanup.cleanup_proven);
assert_eq!(report.target_resolution.profile.id, "qwen3-embedding-0.6b");
}
#[test]
fn swap_to_hash_fallback_is_swapped_to_hash_when_cleanup_proven() {
let from = profile_by_id("qwen3-32b-q4").unwrap();
let to = profile_by_id("tsift-local-hash-v1").unwrap();
let (pre, post) = probe_pair(28_000, 200);
let report = build_swap_report(
from,
to,
pre,
post,
None,
Some(42),
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(report.swap_status, SwapStatus::SwappedToHash);
assert!(report.unload.cleanup.cleanup_proven);
}
#[test]
fn swap_blocks_when_source_unload_not_proven() {
let from = profile_by_id("qwen3-32b-q4").unwrap();
let to = profile_by_id("qwen3-embedding-0.6b").unwrap();
let pre = probe_with_used_vram(200);
let mut post = probe_with_used_vram(8_000);
post.processes.push(GpuProcess {
pid: Some(42),
process_name: "llama-server".to_string(),
used_memory_mib: Some(7_000),
});
let report = build_swap_report(
from,
to,
pre,
post,
None,
Some(42),
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(report.swap_status, SwapStatus::UnloadNotProven);
assert!(!report.unload.cleanup.cleanup_proven);
assert!(
report
.notes
.iter()
.any(|note| note.contains("DO NOT load target"))
);
}
#[test]
fn swap_reports_target_unselectable_when_post_unload_vram_still_high() {
let from = profile_by_id("qwen3-32b-q4").unwrap();
let to = profile_by_id("qwen3-32b-q4").unwrap();
let pre = probe_with_used_vram(29_500);
let post = probe_with_used_vram(29_600);
let report = build_swap_report(
from,
to,
pre,
post,
None,
None,
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(report.swap_status, SwapStatus::NoOpSameProfile);
let from = profile_by_id("qwen3-32b-q4").unwrap();
let to = profile_by_id("qwen3-embedding-8b").unwrap();
let pre = probe_with_used_vram(30_000);
let post = probe_with_used_vram(30_500);
let report = build_swap_report(
from,
to,
pre,
post,
None,
None,
DEFAULT_IDLE_TTL_SECONDS,
DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
);
assert_eq!(
report.swap_status,
SwapStatus::UnloadProvenTargetUnselectable
);
assert!(
report
.notes
.iter()
.any(|note| note.contains("not selectable on the post-unload probe"))
);
}
#[test]
fn rewrite_unload_body_model_replaces_model_field() {
let original = r#"{"model":"qwen3-32b-q4-ollama-default","prompt":"","keep_alive":0}"#;
let rewritten = rewrite_unload_body_model(original, "hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M");
let value: serde_json::Value =
serde_json::from_str(&rewritten).expect("rewritten body is valid JSON");
assert_eq!(
value["model"].as_str(),
Some("hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M")
);
assert_eq!(value["keep_alive"].as_i64(), Some(0));
assert_eq!(value["prompt"].as_str(), Some(""));
}
#[test]
fn rewrite_unload_body_model_preserves_body_when_override_is_empty() {
let original = r#"{"model":"default-tag","keep_alive":0}"#;
let rewritten = rewrite_unload_body_model(original, "");
assert_eq!(rewritten, original);
}
#[test]
fn rewrite_unload_body_model_falls_back_on_invalid_json() {
let original = "not valid json {{{";
let rewritten = rewrite_unload_body_model(original, "any-tag");
assert_eq!(rewritten, original);
}
#[test]
fn normalize_unload_url_appends_generate_path_for_bare_host() {
let url = normalize_unload_url("http://127.0.0.1:11434");
assert_eq!(url, "http://127.0.0.1:11434/api/generate");
}
#[test]
fn normalize_unload_url_idempotent_for_full_generate_url() {
let url = normalize_unload_url("http://127.0.0.1:11434/api/generate");
assert_eq!(url, "http://127.0.0.1:11434/api/generate");
}
#[test]
fn normalize_unload_url_strips_trailing_slash() {
let url = normalize_unload_url("http://127.0.0.1:11434/");
assert_eq!(url, "http://127.0.0.1:11434/api/generate");
}
#[test]
fn prepare_unload_request_returns_none_for_non_api_actions() {
let noop = ProviderUnloadAction {
kind: UnloadActionKind::Noop,
label: "noop".to_string(),
command: None,
http_method: None,
endpoint: None,
body_json: None,
required: false,
};
assert!(prepare_unload_request(&noop, "any-tag").is_none());
}
#[test]
fn prepare_unload_request_applies_model_override_to_body() {
let action = ProviderUnloadAction {
kind: UnloadActionKind::ProviderApi,
label: "ollama keep_alive zero".to_string(),
command: Some(vec!["ollama".to_string(), "stop".to_string()]),
http_method: Some("POST".to_string()),
endpoint: Some("http://127.0.0.1:11434".to_string()),
body_json: Some(
r#"{"model":"profile-default-tag","prompt":"","keep_alive":0}"#.to_string(),
),
required: true,
};
let req = prepare_unload_request(&action, "override-tag")
.expect("ProviderApi action prepares a request");
assert_eq!(req.url, "http://127.0.0.1:11434/api/generate");
assert!(req.body.contains("\"model\":\"override-tag\""));
assert!(!req.body.contains("profile-default-tag"));
assert_eq!(
req.fallback_command,
Some(vec!["ollama".to_string(), "stop".to_string()])
);
}
#[test]
fn prepare_unload_request_synthesizes_body_when_plan_has_none() {
let action = ProviderUnloadAction {
kind: UnloadActionKind::ProviderApi,
label: "synthesized".to_string(),
command: None,
http_method: Some("POST".to_string()),
endpoint: Some("http://host:11434".to_string()),
body_json: None,
required: true,
};
let req = prepare_unload_request(&action, "synth-tag").unwrap();
assert!(req.body.contains("\"model\":\"synth-tag\""));
assert!(req.body.contains("\"keep_alive\":0"));
}
#[test]
fn execute_unload_actions_reports_non_api_as_skipped() {
let actions = vec![ProviderUnloadAction {
kind: UnloadActionKind::Noop,
label: "no GPU unload required".to_string(),
command: None,
http_method: None,
endpoint: None,
body_json: None,
required: false,
}];
let results = execute_unload_actions(&actions, "any-tag");
assert_eq!(results.len(), 1);
assert!(!results[0].executed);
assert_eq!(results[0].outcome, "skipped: non-API action");
}
}