1use anyhow::{Context, Result, bail};
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4use std::fs;
5use std::io::Write;
6use std::path::{Path, PathBuf};
7use std::process::Command;
8use std::time::{SystemTime, UNIX_EPOCH};
9
10pub const DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB: u64 = 4096;
11pub const DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB: u64 = 768;
12pub const DEFAULT_IDLE_TTL_SECONDS: u64 = 0;
13
14pub const DEFAULT_LLAMA_CPP_ENDPOINT: &str = "http://127.0.0.1:8080/models/unload";
19pub const DEFAULT_OLLAMA_ENDPOINT: &str = "http://127.0.0.1:11434/api/generate";
22pub const DEFAULT_VLLM_ENDPOINT: &str = "http://127.0.0.1:8000/sleep";
25
26pub const LLAMA_CPP_ENDPOINT_ENV_VAR: &str = "TSIFT_LLAMA_CPP_ENDPOINT";
28pub const OLLAMA_ENDPOINT_ENV_VAR: &str = "TSIFT_OLLAMA_ENDPOINT";
30pub const VLLM_ENDPOINT_ENV_VAR: &str = "TSIFT_VLLM_ENDPOINT";
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
34pub enum ModelRole {
35 Extract,
36 Embed,
37 Rerank,
38}
39
40#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
41pub enum ProviderKind {
42 LlamaCpp,
43 Ollama,
44 Vllm,
45 HashFallback,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
49pub enum UnloadStrategy {
50 ProcessExit,
51 OllamaKeepAliveZero,
52 LlamaCppRouterUnload,
53 VllmSleep,
54 None,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
58pub enum ConcurrencyClass {
59 ExclusiveLargeGpu,
60 SharedSmallGpu,
61 CpuOrHash,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub enum LeaseMode {
66 Exclusive,
67 Shared,
68 CpuOrHash,
69}
70
71#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
72pub enum UnloadActionKind {
73 ProviderApi,
74 ProcessExit,
75 Sleep,
76 Noop,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
80pub struct ModelProfile {
81 pub id: &'static str,
82 pub label: &'static str,
83 pub provider: ProviderKind,
84 pub model_ref: &'static str,
85 pub quantization: &'static str,
86 pub roles: Vec<ModelRole>,
87 pub context_tokens: u32,
88 pub estimated_weights_mib: u64,
89 pub estimated_kv_mib: u64,
90 pub runtime_margin_mib: u64,
91 pub concurrency: ConcurrencyClass,
92 pub unload_strategy: UnloadStrategy,
93 pub notes: &'static str,
94}
95
96impl ModelProfile {
97 pub fn estimated_total_mib(&self) -> u64 {
98 self.estimated_weights_mib + self.estimated_kv_mib + self.runtime_margin_mib
99 }
100
101 pub fn supports_role(&self, role: &ModelRole) -> bool {
102 self.roles.iter().any(|candidate| candidate == role)
103 }
104}
105
106#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
107pub struct GpuProcess {
108 pub pid: Option<u32>,
109 pub process_name: String,
110 pub used_memory_mib: Option<u64>,
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
114pub struct GpuProbe {
115 pub timestamp_unix_seconds: Option<u64>,
116 pub available: bool,
117 pub gpu_name: Option<String>,
118 pub driver_version: Option<String>,
119 pub total_vram_mib: Option<u64>,
120 pub used_vram_mib: Option<u64>,
121 pub free_vram_mib: Option<u64>,
122 pub processes: Vec<GpuProcess>,
123 pub error: Option<String>,
124}
125
126impl GpuProbe {
127 pub fn unavailable(error: impl Into<String>) -> Self {
128 Self {
129 timestamp_unix_seconds: Some(current_unix_seconds()),
130 available: false,
131 gpu_name: None,
132 driver_version: None,
133 total_vram_mib: None,
134 used_vram_mib: None,
135 free_vram_mib: None,
136 processes: Vec::new(),
137 error: Some(error.into()),
138 }
139 }
140
141 pub fn synthetic_vram(used_vram_mib: u64) -> Self {
142 Self {
143 timestamp_unix_seconds: Some(current_unix_seconds()),
144 available: true,
145 gpu_name: Some("synthetic GPU".to_string()),
146 driver_version: None,
147 total_vram_mib: None,
148 used_vram_mib: Some(used_vram_mib),
149 free_vram_mib: None,
150 processes: Vec::new(),
151 error: None,
152 }
153 }
154}
155
156#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
157pub struct ProfileSelection {
158 pub profile: ModelProfile,
159 pub selectable: bool,
160 pub reason: String,
161}
162
163#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
164pub struct LocalModelStatusReport {
165 pub gpu_probe: GpuProbe,
166 pub extractor_profiles: Vec<ProfileSelection>,
167 pub embedding_profiles: Vec<ProfileSelection>,
168 pub recommended_extractor: Option<String>,
169 pub recommended_embedding: Option<String>,
170 pub notes: Vec<String>,
171}
172
173#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
174pub struct ProviderUnloadAction {
175 pub kind: UnloadActionKind,
176 pub label: String,
177 pub command: Option<Vec<String>>,
178 pub http_method: Option<String>,
179 pub endpoint: Option<String>,
180 pub body_json: Option<String>,
181 pub required: bool,
182}
183
184#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
185pub struct LocalModelLease {
186 pub lease_id: String,
187 pub mode: LeaseMode,
188 pub profile: ModelProfile,
189 pub pre_load_gpu_probe: GpuProbe,
190 pub provider_endpoint: Option<String>,
191 pub provider_pid: Option<u32>,
192 pub idle_ttl_seconds: u64,
193 pub unload_strategy: UnloadStrategy,
194 pub unload_actions: Vec<ProviderUnloadAction>,
195}
196
197#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
198pub enum VramCleanupStatus {
199 Proven,
200 ProvenByExternalAccounting,
201 NotProven,
202 ProbeUnavailable,
203}
204
205#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
206pub struct VramCleanupEvaluation {
207 pub status: VramCleanupStatus,
208 pub cleanup_proven: bool,
209 pub pre_used_mib: Option<u64>,
210 pub post_used_mib: Option<u64>,
211 pub allowed_post_used_mib: Option<u64>,
212 pub used_delta_mib: Option<i64>,
213 pub external_process_delta_mib: u64,
214 pub blocking_processes: Vec<GpuProcess>,
215 pub reason: String,
216}
217
218#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
219pub struct LocalModelLifecycleReport {
220 pub lease: LocalModelLease,
221 pub post_unload_gpu_probe: GpuProbe,
222 pub cleanup: VramCleanupEvaluation,
223 pub notes: Vec<String>,
224}
225
226pub fn default_model_profiles() -> Vec<ModelProfile> {
227 vec![
228 ModelProfile {
229 id: "qwen3-32b-q4",
230 label: "Qwen3-32B 4-bit",
231 provider: ProviderKind::LlamaCpp,
232 model_ref: "Qwen/Qwen3-32B-GGUF",
233 quantization: "q4",
234 roles: vec![ModelRole::Extract],
235 context_tokens: 32_768,
236 estimated_weights_mib: 20_500,
237 estimated_kv_mib: 4_096,
238 runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
239 concurrency: ConcurrencyClass::ExclusiveLargeGpu,
240 unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
241 notes: "default quality extractor/reasoner for a clear RTX 5090",
242 },
243 ModelProfile {
244 id: "qwen3-30b-a3b-instruct-2507-q4",
245 label: "Qwen3-30B-A3B-Instruct-2507 4-bit",
246 provider: ProviderKind::LlamaCpp,
247 model_ref: "Qwen/Qwen3-30B-A3B-Instruct-2507",
248 quantization: "q4",
249 roles: vec![ModelRole::Extract],
250 context_tokens: 262_144,
251 estimated_weights_mib: 19_000,
252 estimated_kv_mib: 4_096,
253 runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
254 concurrency: ConcurrencyClass::ExclusiveLargeGpu,
255 unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
256 notes: "throughput and long-context extractor fallback",
257 },
258 ModelProfile {
259 id: "qwen3-embedding-0.6b",
260 label: "Qwen3-Embedding-0.6B",
261 provider: ProviderKind::LlamaCpp,
262 model_ref: "Qwen/Qwen3-Embedding-0.6B-GGUF",
263 quantization: "q8_or_f16",
264 roles: vec![ModelRole::Embed, ModelRole::Rerank],
265 context_tokens: 32_768,
266 estimated_weights_mib: 1_200,
267 estimated_kv_mib: 512,
268 runtime_margin_mib: 1_024,
269 concurrency: ConcurrencyClass::SharedSmallGpu,
270 unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
271 notes: "default low-pressure embedding companion",
272 },
273 ModelProfile {
274 id: "qwen3-embedding-4b",
275 label: "Qwen3-Embedding-4B",
276 provider: ProviderKind::LlamaCpp,
277 model_ref: "Qwen/Qwen3-Embedding-4B",
278 quantization: "q4_or_q8",
279 roles: vec![ModelRole::Embed, ModelRole::Rerank],
280 context_tokens: 32_768,
281 estimated_weights_mib: 4_200,
282 estimated_kv_mib: 1_024,
283 runtime_margin_mib: 1_024,
284 concurrency: ConcurrencyClass::SharedSmallGpu,
285 unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
286 notes: "higher-quality embedding candidate",
287 },
288 ModelProfile {
289 id: "qwen3-embedding-8b",
290 label: "Qwen3-Embedding-8B",
291 provider: ProviderKind::LlamaCpp,
292 model_ref: "Qwen/Qwen3-Embedding-8B",
293 quantization: "q4_or_q8",
294 roles: vec![ModelRole::Embed, ModelRole::Rerank],
295 context_tokens: 32_768,
296 estimated_weights_mib: 8_200,
297 estimated_kv_mib: 2_048,
298 runtime_margin_mib: 2_048,
299 concurrency: ConcurrencyClass::SharedSmallGpu,
300 unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
301 notes: "benchmark when vector quality matters",
302 },
303 ModelProfile {
304 id: "qwen3.5-35b-a3b-q4",
305 label: "Qwen3.5-35B-A3B 4-bit",
306 provider: ProviderKind::LlamaCpp,
307 model_ref: "Qwen/Qwen3.5-35B-A3B",
308 quantization: "q4",
309 roles: vec![ModelRole::Extract],
310 context_tokens: 128_000,
311 estimated_weights_mib: 24_000,
312 estimated_kv_mib: 8_192,
313 runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
314 concurrency: ConcurrencyClass::ExclusiveLargeGpu,
315 unload_strategy: UnloadStrategy::LlamaCppRouterUnload,
316 notes: "benchmark-only until a reduced-context single-5090 profile is proven",
317 },
318 ModelProfile {
319 id: "tsift-local-hash-v1",
320 label: "tsift local hash fallback",
321 provider: ProviderKind::HashFallback,
322 model_ref: "builtin",
323 quantization: "none",
324 roles: vec![ModelRole::Embed],
325 context_tokens: 0,
326 estimated_weights_mib: 0,
327 estimated_kv_mib: 0,
328 runtime_margin_mib: 0,
329 concurrency: ConcurrencyClass::CpuOrHash,
330 unload_strategy: UnloadStrategy::None,
331 notes: "deterministic fallback for tests and offline runs",
332 },
333 ModelProfile {
339 id: "qwen3-32b-q4-ollama",
340 label: "Qwen3-32B 4-bit (Ollama)",
341 provider: ProviderKind::Ollama,
342 model_ref: "hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M",
343 quantization: "q4",
344 roles: vec![ModelRole::Extract],
345 context_tokens: 32_768,
346 estimated_weights_mib: 20_500,
347 estimated_kv_mib: 4_096,
348 runtime_margin_mib: DEFAULT_DESKTOP_RUNTIME_MARGIN_MIB,
349 concurrency: ConcurrencyClass::ExclusiveLargeGpu,
350 unload_strategy: UnloadStrategy::OllamaKeepAliveZero,
351 notes: "default Ollama-served quality extractor (lazy: keep_alive:0 unloads VRAM)",
352 },
353 ModelProfile {
354 id: "qwen3-embedding-0.6b-ollama",
355 label: "Qwen3-Embedding-0.6B (Ollama)",
356 provider: ProviderKind::Ollama,
357 model_ref: "hf.co/Qwen/Qwen3-Embedding-0.6B-GGUF",
358 quantization: "q8_or_f16",
359 roles: vec![ModelRole::Embed, ModelRole::Rerank],
360 context_tokens: 32_768,
361 estimated_weights_mib: 1_200,
362 estimated_kv_mib: 512,
363 runtime_margin_mib: 1_024,
364 concurrency: ConcurrencyClass::SharedSmallGpu,
365 unload_strategy: UnloadStrategy::OllamaKeepAliveZero,
366 notes: "default low-pressure Ollama-served embedding companion",
367 },
368 ]
369}
370
371pub fn probe_nvidia_smi() -> GpuProbe {
372 let output = Command::new("nvidia-smi")
373 .args([
374 "--query-gpu=name,driver_version,memory.total,memory.used,memory.free",
375 "--format=csv,noheader,nounits",
376 ])
377 .output();
378
379 let output = match output {
380 Ok(output) => output,
381 Err(error) => return GpuProbe::unavailable(format!("nvidia-smi unavailable: {error}")),
382 };
383
384 if !output.status.success() {
385 return GpuProbe::unavailable(format!(
386 "nvidia-smi failed: {}",
387 String::from_utf8_lossy(&output.stderr).trim()
388 ));
389 }
390
391 let stdout = String::from_utf8_lossy(&output.stdout);
392 match parse_gpu_query(stdout.lines().next().unwrap_or_default()) {
393 Ok(mut probe) => {
394 probe.processes = query_nvidia_compute_processes();
395 probe
396 }
397 Err(error) => GpuProbe::unavailable(error.to_string()),
398 }
399}
400
401pub fn build_status_report(probe_gpu: bool) -> LocalModelStatusReport {
402 let gpu_probe = if probe_gpu {
403 probe_nvidia_smi()
404 } else {
405 GpuProbe::unavailable("gpu probe skipped")
406 };
407 build_status_report_with_probe(gpu_probe)
408}
409
410pub fn build_status_report_with_probe(gpu_probe: GpuProbe) -> LocalModelStatusReport {
411 let profiles = default_model_profiles();
412 let extractor_profiles = rank_profiles_for_role(&profiles, &gpu_probe, ModelRole::Extract);
413 let embedding_profiles = rank_profiles_for_role(&profiles, &gpu_probe, ModelRole::Embed);
414 let recommended_extractor = extractor_profiles
415 .iter()
416 .find(|selection| selection.selectable)
417 .map(|selection| selection.profile.id.to_string());
418 let recommended_embedding = embedding_profiles
419 .iter()
420 .find(|selection| selection.selectable)
421 .map(|selection| selection.profile.id.to_string());
422
423 let mut notes = vec![
424 "large 30B/32B extractor profiles are single-lease on one RTX 5090".to_string(),
425 "use provider unload hooks or process exit after each batch to clear VRAM".to_string(),
426 ];
427 if !gpu_probe.available {
428 notes.push("GPU probe unavailable; profile fit is conservative".to_string());
429 }
430
431 LocalModelStatusReport {
432 gpu_probe,
433 extractor_profiles,
434 embedding_profiles,
435 recommended_extractor,
436 recommended_embedding,
437 notes,
438 }
439}
440
441pub fn profile_by_id(profile_id: &str) -> Option<ModelProfile> {
442 default_model_profiles()
443 .into_iter()
444 .find(|profile| profile.id == profile_id)
445}
446
447pub fn lease_mode_for_profile(profile: &ModelProfile) -> LeaseMode {
448 match profile.concurrency {
449 ConcurrencyClass::ExclusiveLargeGpu => LeaseMode::Exclusive,
450 ConcurrencyClass::SharedSmallGpu => LeaseMode::Shared,
451 ConcurrencyClass::CpuOrHash => LeaseMode::CpuOrHash,
452 }
453}
454
455pub fn build_local_model_lease(
456 profile: ModelProfile,
457 pre_load_gpu_probe: GpuProbe,
458 provider_endpoint: Option<String>,
459 provider_pid: Option<u32>,
460 idle_ttl_seconds: u64,
461) -> LocalModelLease {
462 let timestamp = pre_load_gpu_probe
463 .timestamp_unix_seconds
464 .unwrap_or_else(current_unix_seconds);
465 let lease_id = format!("{}-{timestamp}", profile.id);
466 let unload_actions = build_unload_actions(&profile, provider_endpoint.as_deref(), provider_pid);
467
468 LocalModelLease {
469 lease_id,
470 mode: lease_mode_for_profile(&profile),
471 unload_strategy: profile.unload_strategy.clone(),
472 profile,
473 pre_load_gpu_probe,
474 provider_endpoint,
475 provider_pid,
476 idle_ttl_seconds,
477 unload_actions,
478 }
479}
480
481pub fn build_unload_actions(
482 profile: &ModelProfile,
483 provider_endpoint: Option<&str>,
484 provider_pid: Option<u32>,
485) -> Vec<ProviderUnloadAction> {
486 match profile.unload_strategy {
487 UnloadStrategy::LlamaCppRouterUnload => {
488 let endpoint = resolve_provider_endpoint(&profile.unload_strategy, provider_endpoint);
489 let mut actions = vec![ProviderUnloadAction {
490 kind: UnloadActionKind::ProviderApi,
491 label: "llama.cpp router unload".to_string(),
492 command: None,
493 http_method: Some("POST".to_string()),
494 endpoint: Some(endpoint),
495 body_json: Some(format!(r#"{{"model":"{}"}}"#, profile.model_ref)),
496 required: true,
497 }];
498 if let Some(pid) = provider_pid {
499 actions.push(process_exit_action(
500 pid,
501 "terminate llama.cpp worker if unload is not proven",
502 ));
503 }
504 actions
505 }
506 UnloadStrategy::OllamaKeepAliveZero => vec![
507 ProviderUnloadAction {
508 kind: UnloadActionKind::ProviderApi,
509 label: "ollama keep_alive zero".to_string(),
510 command: None,
511 http_method: Some("POST".to_string()),
512 endpoint: Some(resolve_provider_endpoint(
513 &profile.unload_strategy,
514 provider_endpoint,
515 )),
516 body_json: Some(format!(
517 r#"{{"model":"{}","prompt":"","keep_alive":0}}"#,
518 profile.model_ref
519 )),
520 required: true,
521 },
522 ProviderUnloadAction {
523 kind: UnloadActionKind::ProviderApi,
524 label: "ollama stop fallback".to_string(),
525 command: Some(vec![
526 "ollama".to_string(),
527 "stop".to_string(),
528 profile.model_ref.to_string(),
529 ]),
530 http_method: None,
531 endpoint: None,
532 body_json: None,
533 required: false,
534 },
535 ],
536 UnloadStrategy::VllmSleep => {
537 let endpoint = resolve_provider_endpoint(&profile.unload_strategy, provider_endpoint);
538 vec![ProviderUnloadAction {
539 kind: UnloadActionKind::Sleep,
540 label: "vLLM sleep mode".to_string(),
541 command: None,
542 http_method: Some("POST".to_string()),
543 endpoint: Some(endpoint),
544 body_json: None,
545 required: true,
546 }]
547 }
548 UnloadStrategy::ProcessExit => provider_pid
549 .map(|pid| vec![process_exit_action(pid, "terminate isolated model worker")])
550 .unwrap_or_else(|| {
551 vec![ProviderUnloadAction {
552 kind: UnloadActionKind::ProcessExit,
553 label: "terminate isolated model worker".to_string(),
554 command: None,
555 http_method: None,
556 endpoint: None,
557 body_json: None,
558 required: true,
559 }]
560 }),
561 UnloadStrategy::None => vec![ProviderUnloadAction {
562 kind: UnloadActionKind::Noop,
563 label: "no GPU unload required".to_string(),
564 command: None,
565 http_method: None,
566 endpoint: None,
567 body_json: None,
568 required: false,
569 }],
570 }
571}
572
573#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
578pub struct UnloadActionResult {
579 pub label: String,
580 pub executed: bool,
581 pub outcome: String,
582}
583
584#[derive(Debug, Clone, PartialEq, Eq)]
591pub struct PreparedUnloadRequest {
592 pub label: String,
593 pub url: String,
594 pub body: String,
595 pub fallback_command: Option<Vec<String>>,
596}
597
598pub fn rewrite_unload_body_model(body_json: &str, resolved_model_tag: &str) -> String {
604 if resolved_model_tag.is_empty() {
608 return body_json.to_string();
609 }
610 match serde_json::from_str::<serde_json::Value>(body_json) {
611 Ok(mut value) => {
612 if let Some(obj) = value.as_object_mut() {
613 obj.insert(
614 "model".to_string(),
615 serde_json::Value::String(resolved_model_tag.to_string()),
616 );
617 }
618 serde_json::to_string(&value).unwrap_or_else(|_| body_json.to_string())
619 }
620 Err(_) => body_json.to_string(),
621 }
622}
623
624pub fn prepare_unload_request(
627 action: &ProviderUnloadAction,
628 resolved_model_tag: &str,
629) -> Option<PreparedUnloadRequest> {
630 if action.kind != UnloadActionKind::ProviderApi {
631 return None;
632 }
633 let endpoint = action.endpoint.clone().unwrap_or_default();
634 let url = normalize_unload_url(&endpoint);
635 let body = match action.body_json.as_deref() {
636 Some(template) => rewrite_unload_body_model(template, resolved_model_tag),
637 None => format!(
638 r#"{{"model":"{resolved_model_tag}","prompt":"","keep_alive":0}}"#
639 ),
640 };
641 Some(PreparedUnloadRequest {
642 label: action.label.clone(),
643 url,
644 body,
645 fallback_command: action.command.clone(),
646 })
647}
648
649pub fn normalize_unload_url(endpoint: &str) -> String {
653 let trimmed = endpoint.trim_end_matches('/');
654 if trimmed.ends_with("/api/generate") {
655 return trimmed.to_string();
656 }
657 format!("{}/api/generate", trimmed)
658}
659
660pub fn execute_unload_request(req: &PreparedUnloadRequest) -> UnloadActionResult {
665 match post_unload_http(&req.url, &req.body) {
666 Ok(status) => UnloadActionResult {
667 label: req.label.clone(),
668 executed: true,
669 outcome: format!("HTTP {status}"),
670 },
671 Err(err) => {
672 if let Some(cmd) = &req.fallback_command {
673 let _ = std::process::Command::new(&cmd[0])
674 .args(&cmd[1..])
675 .status()
676 .map_err(|e| {
677 eprintln!("tsift-local-model: unload fallback command failed: {e}");
678 });
679 return UnloadActionResult {
680 label: req.label.clone(),
681 executed: true,
682 outcome: format!("POST failed ({err}); ran fallback `{:?}`", cmd),
683 };
684 }
685 UnloadActionResult {
686 label: req.label.clone(),
687 executed: false,
688 outcome: format!("POST failed ({err}); no fallback"),
689 }
690 }
691 }
692}
693
694pub fn unload_model_at(host: &str, model_tag: &str) -> UnloadActionResult {
699 let req = PreparedUnloadRequest {
700 label: format!("ollama keep_alive zero for {model_tag}"),
701 url: normalize_unload_url(host),
702 body: format!(r#"{{"model":"{model_tag}","prompt":"","keep_alive":0}}"#),
703 fallback_command: Some(vec![
704 "ollama".to_string(),
705 "stop".to_string(),
706 model_tag.to_string(),
707 ]),
708 };
709 execute_unload_request(&req)
710}
711
712pub fn execute_unload_actions(
721 actions: &[ProviderUnloadAction],
722 resolved_model_tag: &str,
723) -> Vec<UnloadActionResult> {
724 let mut results = Vec::with_capacity(actions.len());
725 let mut required_succeeded = false;
726 for action in actions {
727 if required_succeeded {
728 results.push(UnloadActionResult {
729 label: action.label.clone(),
730 executed: false,
731 outcome: "skipped: prior required unload succeeded".to_string(),
732 });
733 continue;
734 }
735 let result = match prepare_unload_request(action, resolved_model_tag) {
736 Some(req) => execute_unload_request(&req),
737 None => UnloadActionResult {
738 label: action.label.clone(),
739 executed: false,
740 outcome: "skipped: non-API action".to_string(),
741 },
742 };
743 if result.executed && action.required {
744 required_succeeded = true;
745 }
746 results.push(result);
747 }
748 results
749}
750
751fn post_unload_http(url: &str, body: &str) -> anyhow::Result<String> {
752 use std::time::Duration;
753
754 let payload = serde_json::from_str::<serde_json::Value>(body)
755 .unwrap_or(serde_json::Value::Null);
756 let agent = ureq::Agent::config_builder()
757 .http_status_as_error(false)
758 .timeout_global(Some(Duration::from_secs(30)))
759 .build()
760 .new_agent();
761 let mut response = agent
762 .post(url)
763 .send_json(payload)
764 .with_context(|| format!("posting unload to {url}"))?;
765 let status = response.status();
766 let text = response
767 .body_mut()
768 .read_to_string()
769 .with_context(|| format!("reading unload response (HTTP {status})"))?;
770 if !status.is_success() {
771 bail!("unload HTTP {status}: {}", truncate_str_local(&text, 200));
772 }
773 Ok(format!("{status}"))
774}
775
776fn truncate_str_local(s: &str, max: usize) -> String {
777 if s.len() <= max {
778 s.to_string()
779 } else {
780 format!("{}…", &s[..max])
781 }
782}
783
784pub fn resolve_provider_endpoint(strategy: &UnloadStrategy, explicit: Option<&str>) -> String {
793 if let Some(explicit) = explicit
794 && !explicit.trim().is_empty()
795 {
796 return explicit.to_string();
797 }
798 let (env_var, default): (&str, &str) = match strategy {
799 UnloadStrategy::LlamaCppRouterUnload => {
800 (LLAMA_CPP_ENDPOINT_ENV_VAR, DEFAULT_LLAMA_CPP_ENDPOINT)
801 }
802 UnloadStrategy::OllamaKeepAliveZero => (OLLAMA_ENDPOINT_ENV_VAR, DEFAULT_OLLAMA_ENDPOINT),
803 UnloadStrategy::VllmSleep => (VLLM_ENDPOINT_ENV_VAR, DEFAULT_VLLM_ENDPOINT),
804 UnloadStrategy::ProcessExit | UnloadStrategy::None => return String::new(),
805 };
806 if let Ok(value) = std::env::var(env_var)
807 && !value.trim().is_empty()
808 {
809 return value;
810 }
811 default.to_string()
812}
813
814pub fn build_lifecycle_report(
815 profile: ModelProfile,
816 pre_load_gpu_probe: GpuProbe,
817 post_unload_gpu_probe: GpuProbe,
818 provider_endpoint: Option<String>,
819 provider_pid: Option<u32>,
820 idle_ttl_seconds: u64,
821 tolerance_mib: u64,
822) -> LocalModelLifecycleReport {
823 let lease = build_local_model_lease(
824 profile,
825 pre_load_gpu_probe.clone(),
826 provider_endpoint,
827 provider_pid,
828 idle_ttl_seconds,
829 );
830 let cleanup = evaluate_vram_cleanup(&pre_load_gpu_probe, &post_unload_gpu_probe, tolerance_mib);
831 let mut notes = vec![match lease.mode {
832 LeaseMode::Exclusive => {
833 "large extractor profile requires an exclusive local-model lease".to_string()
834 }
835 LeaseMode::Shared => "small model profile can share GPU when the margin fits".to_string(),
836 LeaseMode::CpuOrHash => "profile does not require GPU VRAM".to_string(),
837 }];
838 if !cleanup.cleanup_proven {
839 notes.push(
840 "future KG runs should fail if cleanup remains unproven after required unload actions"
841 .to_string(),
842 );
843 }
844
845 LocalModelLifecycleReport {
846 lease,
847 post_unload_gpu_probe,
848 cleanup,
849 notes,
850 }
851}
852
853pub fn evaluate_vram_cleanup(
854 pre_load_gpu_probe: &GpuProbe,
855 post_unload_gpu_probe: &GpuProbe,
856 tolerance_mib: u64,
857) -> VramCleanupEvaluation {
858 let pre_used_mib = pre_load_gpu_probe.used_vram_mib;
859 let post_used_mib = post_unload_gpu_probe.used_vram_mib;
860 let allowed_post_used_mib = pre_used_mib.map(|used| used.saturating_add(tolerance_mib));
861 let used_delta_mib = match (pre_used_mib, post_used_mib) {
862 (Some(pre), Some(post)) => Some(post as i64 - pre as i64),
863 _ => None,
864 };
865
866 if !pre_load_gpu_probe.available
867 || !post_unload_gpu_probe.available
868 || pre_used_mib.is_none()
869 || post_used_mib.is_none()
870 {
871 return VramCleanupEvaluation {
872 status: VramCleanupStatus::ProbeUnavailable,
873 cleanup_proven: false,
874 pre_used_mib,
875 post_used_mib,
876 allowed_post_used_mib,
877 used_delta_mib,
878 external_process_delta_mib: 0,
879 blocking_processes: Vec::new(),
880 reason: "pre-load or post-unload GPU probe is unavailable".to_string(),
881 };
882 }
883
884 let pre_used = pre_used_mib.unwrap();
885 let post_used = post_used_mib.unwrap();
886 let allowed = allowed_post_used_mib.unwrap();
887 if post_used <= allowed {
888 return VramCleanupEvaluation {
889 status: VramCleanupStatus::Proven,
890 cleanup_proven: true,
891 pre_used_mib,
892 post_used_mib,
893 allowed_post_used_mib,
894 used_delta_mib,
895 external_process_delta_mib: 0,
896 blocking_processes: Vec::new(),
897 reason: format!(
898 "post-unload VRAM {post_used} MiB is within {tolerance_mib} MiB of baseline {pre_used} MiB"
899 ),
900 };
901 }
902
903 let blocking_processes = post_unload_gpu_probe
904 .processes
905 .iter()
906 .filter(|process| is_tsift_model_process(process))
907 .cloned()
908 .collect::<Vec<_>>();
909 let external_process_delta_mib =
910 external_process_delta_mib(pre_load_gpu_probe, post_unload_gpu_probe);
911
912 if blocking_processes.is_empty()
913 && post_used <= allowed.saturating_add(external_process_delta_mib)
914 {
915 return VramCleanupEvaluation {
916 status: VramCleanupStatus::ProvenByExternalAccounting,
917 cleanup_proven: true,
918 pre_used_mib,
919 post_used_mib,
920 allowed_post_used_mib,
921 used_delta_mib,
922 external_process_delta_mib,
923 blocking_processes,
924 reason: format!(
925 "post-unload VRAM increase is accounted for by {external_process_delta_mib} MiB of non-tsift GPU processes"
926 ),
927 };
928 }
929
930 VramCleanupEvaluation {
931 status: VramCleanupStatus::NotProven,
932 cleanup_proven: false,
933 pre_used_mib,
934 post_used_mib,
935 allowed_post_used_mib,
936 used_delta_mib,
937 external_process_delta_mib,
938 blocking_processes,
939 reason: format!(
940 "post-unload VRAM {post_used} MiB exceeds allowed {allowed} MiB and cleanup is not externally accounted for"
941 ),
942 }
943}
944
945pub fn rank_profiles_for_role(
946 profiles: &[ModelProfile],
947 probe: &GpuProbe,
948 role: ModelRole,
949) -> Vec<ProfileSelection> {
950 profiles
951 .iter()
952 .filter(|profile| profile.supports_role(&role))
953 .map(|profile| selection_for_profile(profile, probe))
954 .collect()
955}
956
957pub fn format_status_human(report: &LocalModelStatusReport) -> String {
958 let mut out = String::new();
959 out.push_str("Local model status\n");
960 if report.gpu_probe.available {
961 out.push_str(&format!(
962 "GPU: {} | VRAM: {} MiB used / {} MiB total ({} MiB free)\n",
963 report.gpu_probe.gpu_name.as_deref().unwrap_or("unknown"),
964 format_optional_u64(report.gpu_probe.used_vram_mib),
965 format_optional_u64(report.gpu_probe.total_vram_mib),
966 format_optional_u64(report.gpu_probe.free_vram_mib)
967 ));
968 } else {
969 out.push_str(&format!(
970 "GPU: unavailable ({})\n",
971 report.gpu_probe.error.as_deref().unwrap_or("unknown error")
972 ));
973 }
974 out.push_str(&format!(
975 "Recommended extractor: {}\n",
976 report
977 .recommended_extractor
978 .as_deref()
979 .unwrap_or("none selectable")
980 ));
981 out.push_str(&format!(
982 "Recommended embedding: {}\n",
983 report
984 .recommended_embedding
985 .as_deref()
986 .unwrap_or("none selectable")
987 ));
988 out.push_str("\nExtractor profiles:\n");
989 for selection in &report.extractor_profiles {
990 out.push_str(&format!(
991 "- {} [{} MiB est]: {} ({})\n",
992 selection.profile.id,
993 selection.profile.estimated_total_mib(),
994 if selection.selectable {
995 "selectable"
996 } else {
997 "blocked"
998 },
999 selection.reason
1000 ));
1001 }
1002 out.push_str("\nEmbedding profiles:\n");
1003 for selection in &report.embedding_profiles {
1004 out.push_str(&format!(
1005 "- {} [{} MiB est]: {} ({})\n",
1006 selection.profile.id,
1007 selection.profile.estimated_total_mib(),
1008 if selection.selectable {
1009 "selectable"
1010 } else {
1011 "blocked"
1012 },
1013 selection.reason
1014 ));
1015 }
1016 out
1017}
1018
1019pub fn format_lifecycle_human(report: &LocalModelLifecycleReport) -> String {
1020 let mut out = String::new();
1021 out.push_str("Local model lifecycle\n");
1022 out.push_str(&format!(
1023 "Profile: {} ({})\n",
1024 report.lease.profile.id, report.lease.profile.label
1025 ));
1026 out.push_str(&format!(
1027 "Lease: {} | mode: {:?} | idle TTL: {}s\n",
1028 report.lease.lease_id, report.lease.mode, report.lease.idle_ttl_seconds
1029 ));
1030 out.push_str(&format!(
1031 "Pre-load VRAM: {} MiB used\n",
1032 format_optional_u64(report.lease.pre_load_gpu_probe.used_vram_mib)
1033 ));
1034 out.push_str(&format!(
1035 "Post-unload VRAM: {} MiB used\n",
1036 format_optional_u64(report.post_unload_gpu_probe.used_vram_mib)
1037 ));
1038 out.push_str(&format!(
1039 "Cleanup: {:?} ({})\n",
1040 report.cleanup.status, report.cleanup.reason
1041 ));
1042 out.push_str("\nRequired unload actions:\n");
1043 for action in &report.lease.unload_actions {
1044 out.push_str(&format!(
1045 "- {}: {}{}\n",
1046 if action.required {
1047 "required"
1048 } else {
1049 "fallback"
1050 },
1051 action.label,
1052 format_action_detail(action)
1053 ));
1054 }
1055 if !report.cleanup.blocking_processes.is_empty() {
1056 out.push_str("\nBlocking GPU processes:\n");
1057 for process in &report.cleanup.blocking_processes {
1058 out.push_str(&format!(
1059 "- pid={} name={} used={} MiB\n",
1060 process
1061 .pid
1062 .map(|pid| pid.to_string())
1063 .unwrap_or_else(|| "unknown".to_string()),
1064 process.process_name,
1065 format_optional_u64(process.used_memory_mib)
1066 ));
1067 }
1068 }
1069 out
1070}
1071
1072fn format_action_detail(action: &ProviderUnloadAction) -> String {
1073 if let Some(command) = &action.command {
1074 return format!(" | command: {}", command.join(" "));
1075 }
1076 if let Some(endpoint) = &action.endpoint {
1077 return format!(
1078 " | {} {}{}",
1079 action.http_method.as_deref().unwrap_or("POST"),
1080 endpoint,
1081 action
1082 .body_json
1083 .as_ref()
1084 .map(|body| format!(" body={body}"))
1085 .unwrap_or_default()
1086 );
1087 }
1088 String::new()
1089}
1090
1091fn selection_for_profile(profile: &ModelProfile, probe: &GpuProbe) -> ProfileSelection {
1092 if profile.concurrency == ConcurrencyClass::CpuOrHash {
1093 return ProfileSelection {
1094 profile: profile.clone(),
1095 selectable: true,
1096 reason: "does not require GPU VRAM".to_string(),
1097 };
1098 }
1099
1100 let Some(free_vram_mib) = probe.free_vram_mib else {
1101 return ProfileSelection {
1102 profile: profile.clone(),
1103 selectable: false,
1104 reason: "free VRAM unknown".to_string(),
1105 };
1106 };
1107
1108 let required = profile.estimated_total_mib();
1109 if required <= free_vram_mib {
1110 ProfileSelection {
1111 profile: profile.clone(),
1112 selectable: true,
1113 reason: format!("estimated {required} MiB fits in {free_vram_mib} MiB free"),
1114 }
1115 } else {
1116 ProfileSelection {
1117 profile: profile.clone(),
1118 selectable: false,
1119 reason: format!("estimated {required} MiB exceeds {free_vram_mib} MiB free"),
1120 }
1121 }
1122}
1123
1124fn process_exit_action(pid: u32, label: &str) -> ProviderUnloadAction {
1125 ProviderUnloadAction {
1126 kind: UnloadActionKind::ProcessExit,
1127 label: label.to_string(),
1128 command: Some(vec![
1129 "kill".to_string(),
1130 "-TERM".to_string(),
1131 pid.to_string(),
1132 ]),
1133 http_method: None,
1134 endpoint: None,
1135 body_json: None,
1136 required: false,
1137 }
1138}
1139
1140fn external_process_delta_mib(
1141 pre_load_gpu_probe: &GpuProbe,
1142 post_unload_gpu_probe: &GpuProbe,
1143) -> u64 {
1144 post_unload_gpu_probe
1145 .processes
1146 .iter()
1147 .filter(|process| !is_tsift_model_process(process))
1148 .map(|process| {
1149 let before = matching_pre_process(pre_load_gpu_probe, process)
1150 .and_then(|pre| pre.used_memory_mib)
1151 .unwrap_or(0);
1152 process.used_memory_mib.unwrap_or(0).saturating_sub(before)
1153 })
1154 .sum()
1155}
1156
1157fn matching_pre_process<'a>(
1158 pre_load_gpu_probe: &'a GpuProbe,
1159 post_process: &GpuProcess,
1160) -> Option<&'a GpuProcess> {
1161 if let Some(pid) = post_process.pid
1162 && let Some(process) = pre_load_gpu_probe
1163 .processes
1164 .iter()
1165 .find(|candidate| candidate.pid == Some(pid))
1166 {
1167 return Some(process);
1168 }
1169 pre_load_gpu_probe
1170 .processes
1171 .iter()
1172 .find(|candidate| candidate.process_name == post_process.process_name)
1173}
1174
1175fn is_tsift_model_process(process: &GpuProcess) -> bool {
1176 let name = process.process_name.to_ascii_lowercase();
1177 name.contains("tsift")
1178 || name.contains("llama")
1179 || name.contains("ollama")
1180 || name.contains("vllm")
1181 || name.contains("ggml")
1182}
1183
1184pub fn current_unix_seconds() -> u64 {
1185 SystemTime::now()
1186 .duration_since(UNIX_EPOCH)
1187 .map(|duration| duration.as_secs())
1188 .unwrap_or(0)
1189}
1190
1191fn parse_gpu_query(line: &str) -> Result<GpuProbe> {
1192 let parts = line.split(',').map(str::trim).collect::<Vec<_>>();
1193 if parts.len() != 5 {
1194 anyhow::bail!("unexpected nvidia-smi gpu query row: {line}");
1195 }
1196 Ok(GpuProbe {
1197 timestamp_unix_seconds: Some(current_unix_seconds()),
1198 available: true,
1199 gpu_name: Some(parts[0].to_string()),
1200 driver_version: Some(parts[1].to_string()),
1201 total_vram_mib: Some(parse_optional_u64(parts[2]).context("parse total VRAM")?),
1202 used_vram_mib: Some(parse_optional_u64(parts[3]).context("parse used VRAM")?),
1203 free_vram_mib: Some(parse_optional_u64(parts[4]).context("parse free VRAM")?),
1204 processes: Vec::new(),
1205 error: None,
1206 })
1207}
1208
1209fn query_nvidia_compute_processes() -> Vec<GpuProcess> {
1210 let Ok(output) = Command::new("nvidia-smi")
1211 .args([
1212 "--query-compute-apps=pid,process_name,used_memory",
1213 "--format=csv,noheader,nounits",
1214 ])
1215 .output()
1216 else {
1217 return Vec::new();
1218 };
1219 if !output.status.success() {
1220 return Vec::new();
1221 }
1222 String::from_utf8_lossy(&output.stdout)
1223 .lines()
1224 .filter_map(parse_process_query)
1225 .collect()
1226}
1227
1228fn parse_process_query(line: &str) -> Option<GpuProcess> {
1229 let parts = line.split(',').map(str::trim).collect::<Vec<_>>();
1230 if parts.len() != 3 || parts.iter().all(|part| part.is_empty()) {
1231 return None;
1232 }
1233 Some(GpuProcess {
1234 pid: parts[0].parse::<u32>().ok(),
1235 process_name: parts[1].to_string(),
1236 used_memory_mib: parse_optional_u64(parts[2]).ok(),
1237 })
1238}
1239
1240fn parse_optional_u64(input: &str) -> Result<u64> {
1241 let cleaned = input.trim().trim_end_matches("MiB").trim();
1242 cleaned
1243 .parse::<u64>()
1244 .with_context(|| format!("parse integer from {input:?}"))
1245}
1246
1247fn format_optional_u64(value: Option<u64>) -> String {
1248 value
1249 .map(|value| value.to_string())
1250 .unwrap_or_else(|| "unknown".to_string())
1251}
1252
1253pub const LEASE_REGISTRY_VERSION: u32 = 1;
1269pub const DEFAULT_LEASE_TTL_SECONDS: u64 = 0;
1271pub const LEASE_FILE_ENV_VAR: &str = "TSIFT_LEASE_FILE";
1273
1274#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1276pub struct GpuLeaseRecord {
1277 pub profile_id: String,
1278 pub holder_pid: u32,
1279 pub holder_command: String,
1280 pub acquired_at_unix_seconds: u64,
1285 pub lease_mode: LeaseMode,
1286 pub vram_baseline_mib: u64,
1287 pub idle_ttl_seconds: u64,
1288 pub notes: Vec<String>,
1289}
1290
1291#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1293pub struct GpuLeaseRegistry {
1294 pub version: u32,
1295 pub leases: BTreeMap<String, Vec<GpuLeaseRecord>>,
1296}
1297
1298impl Default for GpuLeaseRegistry {
1299 fn default() -> Self {
1300 Self {
1301 version: LEASE_REGISTRY_VERSION,
1302 leases: BTreeMap::new(),
1303 }
1304 }
1305}
1306
1307#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1308pub enum GpuLeaseAcquisitionStatus {
1309 Acquired,
1311 Refreshed,
1313 ReclaimedStale,
1315 CpuOrHashBypass,
1317 Conflict,
1319}
1320
1321#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1322pub struct GpuLeaseConflict {
1323 pub profile_id: String,
1324 pub holder_pid: u32,
1325 pub holder_command: String,
1326 pub acquired_at_unix_seconds: u64,
1327 pub lease_mode: LeaseMode,
1328}
1329
1330#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1331pub struct GpuLeaseAcquisition {
1332 pub profile_id: String,
1333 pub holder_pid: u32,
1334 pub status: GpuLeaseAcquisitionStatus,
1335 pub record: Option<GpuLeaseRecord>,
1336 pub conflict: Option<GpuLeaseConflict>,
1337 pub reclaimed: Vec<GpuLeaseRecord>,
1339}
1340
1341#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1342pub enum GpuLeaseReleaseOutcome {
1343 Released,
1345 NotHeld,
1347 ProfileAbsent,
1349}
1350
1351#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1352pub struct GpuLeaseRelease {
1353 pub profile_id: String,
1354 pub holder_pid: u32,
1355 pub outcome: GpuLeaseReleaseOutcome,
1356 pub remaining_holders: u32,
1358}
1359
1360pub fn resolve_lease_file(override_path: Option<&Path>) -> PathBuf {
1366 if let Some(path) = override_path {
1367 return path.to_path_buf();
1368 }
1369 if let Ok(path) = std::env::var(LEASE_FILE_ENV_VAR) {
1370 return PathBuf::from(path);
1371 }
1372 if let Ok(state_dir) = std::env::var("XDG_STATE_HOME")
1373 && !state_dir.is_empty()
1374 {
1375 return PathBuf::from(state_dir)
1376 .join("tsift")
1377 .join("gpu-lease.json");
1378 }
1379 if let Ok(home) = std::env::var("HOME")
1380 && !home.is_empty()
1381 {
1382 return PathBuf::from(home).join(".tsift").join("gpu-lease.json");
1383 }
1384 PathBuf::from("./.tsift/gpu-lease.json")
1385}
1386
1387pub fn is_pid_alive(pid: u32) -> bool {
1393 if pid == 0 {
1394 return false;
1395 }
1396 if pid == std::process::id() {
1397 return true;
1398 }
1399 match Command::new("kill").arg("-0").arg(pid.to_string()).output() {
1400 Ok(output) => output.status.success(),
1401 Err(_) => false,
1402 }
1403}
1404
1405pub fn registry_lock_path(path: &Path) -> PathBuf {
1411 let mut name = path.as_os_str().to_os_string();
1412 name.push(".lock");
1413 PathBuf::from(name)
1414}
1415
1416fn with_registry_lock<T>(path: &Path, op: impl FnOnce() -> Result<T>) -> Result<T> {
1427 use fs4::fs_std::FileExt;
1428 let lock_path = registry_lock_path(path);
1429 if let Some(parent) = lock_path.parent()
1430 && !parent.as_os_str().is_empty()
1431 {
1432 fs::create_dir_all(parent).context("create lease registry lock parent")?;
1433 }
1434 let lock_file = fs::OpenOptions::new()
1435 .create(true)
1436 .read(true)
1437 .write(true)
1438 .truncate(false)
1439 .open(&lock_path)
1440 .with_context(|| format!("open lease registry lock {}", lock_path.display()))?;
1441 lock_file
1442 .lock_exclusive()
1443 .context("acquire exclusive lease registry lock")?;
1444 let result = op();
1445 let _ = FileExt::unlock(&lock_file);
1448 result
1449}
1450
1451pub fn read_lease_registry(path: &Path) -> Result<GpuLeaseRegistry> {
1453 match fs::read_to_string(path) {
1454 Ok(contents) => {
1455 if contents.trim().is_empty() {
1456 return Ok(GpuLeaseRegistry::default());
1457 }
1458 serde_json::from_str(&contents).context("parse gpu lease registry")
1459 }
1460 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
1461 Ok(GpuLeaseRegistry::default())
1462 }
1463 Err(error) => Err(error).context("read gpu lease registry"),
1464 }
1465}
1466
1467pub fn write_lease_registry(path: &Path, registry: &GpuLeaseRegistry) -> Result<()> {
1469 if let Some(parent) = path.parent()
1470 && !parent.as_os_str().is_empty()
1471 {
1472 fs::create_dir_all(parent).context("create lease registry parent")?;
1473 }
1474 let payload = serde_json::to_string_pretty(registry).context("serialize lease registry")?;
1475 let temp_path = path.with_extension(format!(
1476 "json.tmp.{}.{}",
1477 std::process::id(),
1478 current_unix_seconds()
1479 ));
1480 let mut handle = fs::File::create(&temp_path).context("create lease registry temp file")?;
1481 handle
1482 .write_all(payload.as_bytes())
1483 .context("write lease registry temp file")?;
1484 handle.sync_all().context("sync lease registry temp file")?;
1485 drop(handle);
1486 fs::rename(&temp_path, path).context("rename lease registry into place")?;
1487 Ok(())
1488}
1489
1490pub fn prune_stale_leases(
1497 registry: &mut GpuLeaseRegistry,
1498 now: u64,
1499 is_alive: impl Fn(u32) -> bool,
1500) -> Vec<GpuLeaseRecord> {
1501 let mut pruned = Vec::new();
1502 let mut empty_keys = Vec::new();
1503 for (profile_id, holders) in registry.leases.iter_mut() {
1504 let mut kept = Vec::with_capacity(holders.len());
1505 for record in holders.drain(..) {
1506 let pid_dead = !is_alive(record.holder_pid);
1507 let ttl_expired = record.idle_ttl_seconds > 0
1508 && now.saturating_sub(record.acquired_at_unix_seconds) > record.idle_ttl_seconds;
1509 if pid_dead || ttl_expired {
1510 pruned.push(record);
1511 } else {
1512 kept.push(record);
1513 }
1514 }
1515 if kept.is_empty() {
1516 empty_keys.push(profile_id.clone());
1517 }
1518 *holders = kept;
1519 }
1520 for key in empty_keys {
1521 registry.leases.remove(&key);
1522 }
1523 pruned
1524}
1525
1526#[allow(clippy::too_many_arguments)]
1531pub fn apply_acquire(
1532 registry: &mut GpuLeaseRegistry,
1533 profile: &ModelProfile,
1534 holder_pid: u32,
1535 holder_command: &str,
1536 vram_baseline_mib: u64,
1537 idle_ttl_seconds: u64,
1538 now: u64,
1539 is_alive: impl Fn(u32) -> bool,
1540) -> GpuLeaseAcquisition {
1541 if profile.concurrency == ConcurrencyClass::CpuOrHash {
1542 return GpuLeaseAcquisition {
1543 profile_id: profile.id.to_string(),
1544 holder_pid,
1545 status: GpuLeaseAcquisitionStatus::CpuOrHashBypass,
1546 record: None,
1547 conflict: None,
1548 reclaimed: Vec::new(),
1549 };
1550 }
1551
1552 let reclaimed = prune_stale_leases(registry, now, &is_alive);
1553 let mode = lease_mode_for_profile(profile);
1554 let entry = registry.leases.entry(profile.id.to_string()).or_default();
1555 let already_held = entry
1556 .iter()
1557 .position(|record| record.holder_pid == holder_pid);
1558
1559 let record = GpuLeaseRecord {
1560 profile_id: profile.id.to_string(),
1561 holder_pid,
1562 holder_command: holder_command.to_string(),
1563 acquired_at_unix_seconds: now,
1564 lease_mode: mode.clone(),
1565 vram_baseline_mib,
1566 idle_ttl_seconds,
1567 notes: Vec::new(),
1568 };
1569
1570 let status = if let Some(index) = already_held {
1571 entry[index] = record.clone();
1572 GpuLeaseAcquisitionStatus::Refreshed
1573 } else {
1574 match mode {
1575 LeaseMode::Exclusive => {
1576 if let Some(blocker) = entry.first() {
1577 return GpuLeaseAcquisition {
1578 profile_id: profile.id.to_string(),
1579 holder_pid,
1580 status: GpuLeaseAcquisitionStatus::Conflict,
1581 record: None,
1582 conflict: Some(GpuLeaseConflict {
1583 profile_id: profile.id.to_string(),
1584 holder_pid: blocker.holder_pid,
1585 holder_command: blocker.holder_command.clone(),
1586 acquired_at_unix_seconds: blocker.acquired_at_unix_seconds,
1587 lease_mode: blocker.lease_mode.clone(),
1588 }),
1589 reclaimed,
1590 };
1591 }
1592 entry.push(record.clone());
1593 if reclaimed
1594 .iter()
1595 .any(|pruned| pruned.profile_id == profile.id)
1596 {
1597 GpuLeaseAcquisitionStatus::ReclaimedStale
1598 } else {
1599 GpuLeaseAcquisitionStatus::Acquired
1600 }
1601 }
1602 LeaseMode::Shared => {
1603 entry.push(record.clone());
1604 if reclaimed
1605 .iter()
1606 .any(|pruned| pruned.profile_id == profile.id)
1607 {
1608 GpuLeaseAcquisitionStatus::ReclaimedStale
1609 } else {
1610 GpuLeaseAcquisitionStatus::Acquired
1611 }
1612 }
1613 LeaseMode::CpuOrHash => GpuLeaseAcquisitionStatus::CpuOrHashBypass,
1614 }
1615 };
1616
1617 GpuLeaseAcquisition {
1618 profile_id: profile.id.to_string(),
1619 holder_pid,
1620 status,
1621 record: Some(record),
1622 conflict: None,
1623 reclaimed,
1624 }
1625}
1626
1627pub fn apply_release(
1629 registry: &mut GpuLeaseRegistry,
1630 profile_id: &str,
1631 holder_pid: u32,
1632 now: u64,
1633 is_alive: impl Fn(u32) -> bool,
1634) -> GpuLeaseRelease {
1635 let _ = prune_stale_leases(registry, now, &is_alive);
1636 let Some(holders) = registry.leases.get_mut(profile_id) else {
1637 return GpuLeaseRelease {
1638 profile_id: profile_id.to_string(),
1639 holder_pid,
1640 outcome: GpuLeaseReleaseOutcome::ProfileAbsent,
1641 remaining_holders: 0,
1642 };
1643 };
1644 let before = holders.len();
1645 holders.retain(|record| record.holder_pid != holder_pid);
1646 let removed = before - holders.len();
1647 let remaining = holders.len() as u32;
1648 if holders.is_empty() {
1649 registry.leases.remove(profile_id);
1651 }
1652 let outcome = if removed == 0 {
1653 GpuLeaseReleaseOutcome::NotHeld
1654 } else {
1655 GpuLeaseReleaseOutcome::Released
1656 };
1657 GpuLeaseRelease {
1658 profile_id: profile_id.to_string(),
1659 holder_pid,
1660 outcome,
1661 remaining_holders: remaining,
1662 }
1663}
1664
1665pub fn acquire_lease(
1667 profile_id: &str,
1668 holder_pid: u32,
1669 holder_command: &str,
1670 vram_baseline_mib: u64,
1671 idle_ttl_seconds: u64,
1672 now: u64,
1673 path: &Path,
1674) -> Result<GpuLeaseAcquisition> {
1675 let profile = profile_by_id(profile_id)
1676 .with_context(|| format!("unknown local model profile {profile_id:?}"))?;
1677 with_registry_lock(path, || {
1678 let mut registry = read_lease_registry(path)?;
1679 let acquisition = apply_acquire(
1680 &mut registry,
1681 &profile,
1682 holder_pid,
1683 holder_command,
1684 vram_baseline_mib,
1685 idle_ttl_seconds,
1686 now,
1687 is_pid_alive,
1688 );
1689 if acquisition.status != GpuLeaseAcquisitionStatus::CpuOrHashBypass {
1691 write_lease_registry(path, ®istry)?;
1692 }
1693 Ok(acquisition)
1694 })
1695}
1696
1697pub fn release_lease(
1699 profile_id: &str,
1700 holder_pid: u32,
1701 now: u64,
1702 path: &Path,
1703) -> Result<GpuLeaseRelease> {
1704 with_registry_lock(path, || {
1705 let mut registry = read_lease_registry(path)?;
1706 let release = apply_release(&mut registry, profile_id, holder_pid, now, is_pid_alive);
1707 write_lease_registry(path, ®istry)?;
1708 Ok(release)
1709 })
1710}
1711
1712pub fn show_registry(path: &Path, now: u64, include_stale: bool) -> Result<GpuLeaseRegistry> {
1715 let mut registry = read_lease_registry(path)?;
1716 if !include_stale {
1717 prune_stale_leases(&mut registry, now, is_pid_alive);
1718 }
1719 Ok(registry)
1720}
1721
1722#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1723pub enum GpuLeaseRenewOutcome {
1724 Renewed,
1726 NotHeld,
1728 ProfileAbsent,
1730}
1731
1732#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1733pub struct GpuLeaseRenew {
1734 pub profile_id: String,
1735 pub holder_pid: u32,
1736 pub outcome: GpuLeaseRenewOutcome,
1737 pub renewed_at_unix_seconds: Option<u64>,
1739}
1740
1741pub fn apply_renew(
1744 registry: &mut GpuLeaseRegistry,
1745 profile_id: &str,
1746 holder_pid: u32,
1747 now: u64,
1748 is_alive: impl Fn(u32) -> bool,
1749) -> GpuLeaseRenew {
1750 let _ = prune_stale_leases(registry, now, &is_alive);
1751 let Some(holders) = registry.leases.get_mut(profile_id) else {
1752 return GpuLeaseRenew {
1753 profile_id: profile_id.to_string(),
1754 holder_pid,
1755 outcome: GpuLeaseRenewOutcome::ProfileAbsent,
1756 renewed_at_unix_seconds: None,
1757 };
1758 };
1759 if let Some(record) = holders
1760 .iter_mut()
1761 .find(|record| record.holder_pid == holder_pid)
1762 {
1763 record.acquired_at_unix_seconds = now;
1764 GpuLeaseRenew {
1765 profile_id: profile_id.to_string(),
1766 holder_pid,
1767 outcome: GpuLeaseRenewOutcome::Renewed,
1768 renewed_at_unix_seconds: Some(now),
1769 }
1770 } else {
1771 GpuLeaseRenew {
1772 profile_id: profile_id.to_string(),
1773 holder_pid,
1774 outcome: GpuLeaseRenewOutcome::NotHeld,
1775 renewed_at_unix_seconds: None,
1776 }
1777 }
1778}
1779
1780pub fn renew_lease(
1785 profile_id: &str,
1786 holder_pid: u32,
1787 now: u64,
1788 path: &Path,
1789) -> Result<GpuLeaseRenew> {
1790 with_registry_lock(path, || {
1791 let mut registry = read_lease_registry(path)?;
1792 let renew = apply_renew(&mut registry, profile_id, holder_pid, now, is_pid_alive);
1793 write_lease_registry(path, ®istry)?;
1794 Ok(renew)
1795 })
1796}
1797
1798#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1799pub struct GpuLeaseReap {
1800 pub reclaimed: Vec<GpuLeaseRecord>,
1802 pub emptied_profiles: Vec<String>,
1805}
1806
1807pub fn reap_leases(now: u64, path: &Path) -> Result<GpuLeaseReap> {
1815 with_registry_lock(path, || {
1816 let mut registry = read_lease_registry(path)?;
1817 let before: std::collections::BTreeSet<String> =
1818 registry.leases.keys().cloned().collect();
1819 let reclaimed = prune_stale_leases(&mut registry, now, is_pid_alive);
1820 let emptied_profiles: Vec<String> = before
1821 .into_iter()
1822 .filter(|profile_id| !registry.leases.contains_key(profile_id))
1823 .collect();
1824 write_lease_registry(path, ®istry)?;
1825 Ok(GpuLeaseReap {
1826 reclaimed,
1827 emptied_profiles,
1828 })
1829 })
1830}
1831
1832pub fn format_lease_show_human(registry: &GpuLeaseRegistry, now: u64) -> String {
1834 let mut out = String::new();
1835 out.push_str("GPU lease registry\n");
1836 out.push_str(&format!("version: {}\n", registry.version));
1837 if registry.leases.is_empty() {
1838 out.push_str("leases: none\n");
1839 return out;
1840 }
1841 out.push_str(&format!("profiles held: {}\n", registry.leases.len()));
1842 for (profile_id, holders) in ®istry.leases {
1843 out.push_str(&format!("\n{profile_id} ({} holder(s)):\n", holders.len()));
1844 for record in holders {
1845 let age = now.saturating_sub(record.acquired_at_unix_seconds);
1846 out.push_str(&format!(
1847 " pid={} cmd={} mode={:?} acquired={}s ago baseline={} MiB ttl={}s",
1848 record.holder_pid,
1849 record.holder_command,
1850 record.lease_mode,
1851 age,
1852 record.vram_baseline_mib,
1853 record.idle_ttl_seconds
1854 ));
1855 if record.notes.is_empty() {
1856 out.push('\n');
1857 } else {
1858 out.push_str(&format!(" notes={}\n", record.notes.join("; ")));
1859 }
1860 }
1861 }
1862 out
1863}
1864
1865#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1880#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
1881pub enum ProfilePreference {
1882 Auto,
1884 Pinned(String),
1887 ForceHash,
1890}
1891
1892impl ProfilePreference {
1893 pub fn from_cli(value: Option<&str>) -> Self {
1898 match value.map(str::trim) {
1899 None | Some("") => ProfilePreference::Auto,
1900 Some("hash") | Some("tsift-local-hash-v1") => ProfilePreference::ForceHash,
1901 Some(other) => ProfilePreference::Pinned(other.to_string()),
1902 }
1903 }
1904
1905 pub fn describe(&self) -> String {
1906 match self {
1907 ProfilePreference::Auto => "auto".to_string(),
1908 ProfilePreference::Pinned(id) => format!("pinned:{id}"),
1909 ProfilePreference::ForceHash => "force-hash".to_string(),
1910 }
1911 }
1912}
1913
1914#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1916#[serde(rename_all = "snake_case")]
1917pub enum ProfileResolutionSource {
1918 AutoRanked,
1920 Pinned,
1922 PinnedUnselectable,
1926 ForcedHash,
1928}
1929
1930#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1932pub struct ProfileResolution {
1933 pub preference: ProfilePreference,
1934 pub role: ModelRole,
1935 pub source: ProfileResolutionSource,
1936 pub profile: ModelProfile,
1937 pub selectable: bool,
1938 pub reason: String,
1939}
1940
1941pub fn resolve_profile_preference(
1947 preference: &ProfilePreference,
1948 role: ModelRole,
1949 probe: &GpuProbe,
1950) -> ProfileResolution {
1951 let profiles = default_model_profiles();
1952 let hash_profile = profiles
1953 .iter()
1954 .find(|profile| profile.id == "tsift-local-hash-v1")
1955 .cloned()
1956 .expect("hash fallback profile is always present");
1957
1958 match preference {
1959 ProfilePreference::ForceHash => ProfileResolution {
1960 preference: preference.clone(),
1961 role,
1962 source: ProfileResolutionSource::ForcedHash,
1963 profile: hash_profile,
1964 selectable: true,
1965 reason: "caller forced the CPU/hash fallback".to_string(),
1966 },
1967 ProfilePreference::Auto => {
1968 let ranked = rank_profiles_for_role(&profiles, probe, role);
1969 let pick = ranked
1970 .iter()
1971 .find(|selection| selection.selectable)
1972 .cloned()
1973 .or_else(|| {
1974 ranked.into_iter().next().map(|selection| ProfileSelection {
1975 selectable: false,
1976 ..selection
1977 })
1978 });
1979 match pick {
1980 Some(selection) if selection.selectable => ProfileResolution {
1981 preference: preference.clone(),
1982 role,
1983 source: ProfileResolutionSource::AutoRanked,
1984 profile: selection.profile.clone(),
1985 selectable: true,
1986 reason: format!("auto-ranked: {}", selection.reason),
1987 },
1988 Some(selection) => ProfileResolution {
1989 preference: preference.clone(),
1990 role,
1991 source: ProfileResolutionSource::AutoRanked,
1992 profile: hash_profile,
1993 selectable: true,
1994 reason: format!(
1995 "auto-ranked but no GPU profile selectable ({}); using hash fallback",
1996 selection.reason
1997 ),
1998 },
1999 None => ProfileResolution {
2000 preference: preference.clone(),
2001 role,
2002 source: ProfileResolutionSource::AutoRanked,
2003 profile: hash_profile,
2004 selectable: true,
2005 reason: "no profile matches the requested role; using hash fallback"
2006 .to_string(),
2007 },
2008 }
2009 }
2010 ProfilePreference::Pinned(id) => match profile_by_id(id) {
2011 Some(profile) if profile.supports_role(&role) => {
2012 let selection = selection_for_profile(&profile, probe);
2013 if selection.selectable {
2014 ProfileResolution {
2015 preference: preference.clone(),
2016 role,
2017 source: ProfileResolutionSource::Pinned,
2018 profile,
2019 selectable: true,
2020 reason: format!("pinned: {}", selection.reason),
2021 }
2022 } else {
2023 ProfileResolution {
2024 preference: preference.clone(),
2025 role,
2026 source: ProfileResolutionSource::PinnedUnselectable,
2027 profile: hash_profile,
2028 selectable: true,
2029 reason: format!(
2030 "pinned {} is not selectable ({}); using hash fallback",
2031 id, selection.reason
2032 ),
2033 }
2034 }
2035 }
2036 Some(_) => ProfileResolution {
2037 preference: preference.clone(),
2038 role,
2039 source: ProfileResolutionSource::PinnedUnselectable,
2040 profile: hash_profile,
2041 selectable: true,
2042 reason: format!(
2043 "pinned {id} does not support role {:?}; using hash fallback",
2044 role
2045 ),
2046 },
2047 None => ProfileResolution {
2048 preference: preference.clone(),
2049 role,
2050 source: ProfileResolutionSource::PinnedUnselectable,
2051 profile: hash_profile,
2052 selectable: true,
2053 reason: format!("pinned profile id {id:?} is unknown; using hash fallback"),
2054 },
2055 },
2056 }
2057}
2058
2059#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2073#[serde(rename_all = "snake_case")]
2074pub enum SwapStatus {
2075 Swapped,
2077 SwappedToHash,
2080 UnloadProvenTargetUnselectable,
2083 UnloadNotProven,
2086 NoOpSameProfile,
2088}
2089
2090#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2091pub struct LocalModelSwapReport {
2092 pub from_profile_id: String,
2093 pub to_profile_id: String,
2094 pub unload: LocalModelLifecycleReport,
2095 pub target_resolution: ProfileResolution,
2096 pub swap_status: SwapStatus,
2097 pub notes: Vec<String>,
2098}
2099
2100#[allow(clippy::too_many_arguments)]
2106pub fn build_swap_report(
2107 from_profile: ModelProfile,
2108 to_profile: ModelProfile,
2109 pre_load_probe: GpuProbe,
2110 post_unload_probe: GpuProbe,
2111 provider_endpoint: Option<String>,
2112 provider_pid: Option<u32>,
2113 idle_ttl_seconds: u64,
2114 tolerance_mib: u64,
2115) -> LocalModelSwapReport {
2116 let unload = build_lifecycle_report(
2117 from_profile.clone(),
2118 pre_load_probe,
2119 post_unload_probe.clone(),
2120 provider_endpoint,
2121 provider_pid,
2122 idle_ttl_seconds,
2123 tolerance_mib,
2124 );
2125
2126 let target_role = to_profile
2127 .roles
2128 .first()
2129 .copied()
2130 .unwrap_or(ModelRole::Extract);
2131 let target_resolution = resolve_profile_preference(
2132 &ProfilePreference::Pinned(to_profile.id.to_string()),
2133 target_role,
2134 &post_unload_probe,
2135 );
2136
2137 let swap_status = if from_profile.id == to_profile.id {
2138 SwapStatus::NoOpSameProfile
2139 } else if !unload.cleanup.cleanup_proven {
2140 SwapStatus::UnloadNotProven
2141 } else if to_profile.concurrency == ConcurrencyClass::CpuOrHash {
2142 SwapStatus::SwappedToHash
2143 } else if target_resolution.selectable && target_resolution.profile.id == to_profile.id {
2144 SwapStatus::Swapped
2145 } else {
2146 SwapStatus::UnloadProvenTargetUnselectable
2147 };
2148
2149 let mut notes = vec![
2150 format!("swapping from {} to {}", from_profile.id, to_profile.id),
2151 format!("unload cleanup: {:?}", unload.cleanup.status),
2152 format!("target resolution: {:?}", target_resolution.source),
2153 ];
2154 if swap_status == SwapStatus::UnloadNotProven {
2155 notes.push("DO NOT load target — source unload did not prove VRAM cleanup".to_string());
2156 }
2157 if swap_status == SwapStatus::UnloadProvenTargetUnselectable {
2158 notes.push(format!(
2159 "target {} is not selectable on the post-unload probe; consider the hash fallback or a smaller profile",
2160 to_profile.id
2161 ));
2162 }
2163
2164 LocalModelSwapReport {
2165 from_profile_id: from_profile.id.to_string(),
2166 to_profile_id: to_profile.id.to_string(),
2167 unload,
2168 target_resolution,
2169 swap_status,
2170 notes,
2171 }
2172}
2173
2174#[cfg(test)]
2175mod tests {
2176 use super::*;
2177
2178 fn rtx_5090_probe() -> GpuProbe {
2179 GpuProbe {
2180 timestamp_unix_seconds: Some(1_781_000_000),
2181 available: true,
2182 gpu_name: Some("NVIDIA GeForce RTX 5090".to_string()),
2183 driver_version: Some("610.43.02".to_string()),
2184 total_vram_mib: Some(32_607),
2185 used_vram_mib: Some(179),
2186 free_vram_mib: Some(32_428),
2187 processes: Vec::new(),
2188 error: None,
2189 }
2190 }
2191
2192 fn probe_with_used_vram(used_vram_mib: u64) -> GpuProbe {
2193 let mut probe = rtx_5090_probe();
2194 probe.used_vram_mib = Some(used_vram_mib);
2195 probe.free_vram_mib = probe
2196 .total_vram_mib
2197 .map(|total| total.saturating_sub(used_vram_mib));
2198 probe
2199 }
2200
2201 #[test]
2202 fn qwen3_32b_is_default_extractor_for_clear_5090() {
2203 let report = build_status_report_with_probe(rtx_5090_probe());
2204 assert_eq!(
2205 report.recommended_extractor.as_deref(),
2206 Some("qwen3-32b-q4")
2207 );
2208 assert!(
2209 report
2210 .extractor_profiles
2211 .iter()
2212 .any(|selection| selection.profile.id == "qwen3.5-35b-a3b-q4"
2213 && !selection.selectable)
2214 );
2215 }
2216
2217 #[test]
2218 fn hash_fallback_selects_without_gpu_probe() {
2219 let report = build_status_report_with_probe(GpuProbe::unavailable("missing"));
2220 assert_eq!(
2221 report.recommended_embedding.as_deref(),
2222 Some("tsift-local-hash-v1")
2223 );
2224 assert_eq!(report.recommended_extractor, None);
2225 }
2226
2227 #[test]
2228 fn parses_nvidia_smi_gpu_query_row() {
2229 let probe =
2230 parse_gpu_query("NVIDIA GeForce RTX 5090, 610.43.02, 32607, 179, 32428").unwrap();
2231 assert!(probe.timestamp_unix_seconds.is_some());
2232 assert_eq!(probe.gpu_name.as_deref(), Some("NVIDIA GeForce RTX 5090"));
2233 assert_eq!(probe.total_vram_mib, Some(32_607));
2234 assert_eq!(probe.used_vram_mib, Some(179));
2235 assert_eq!(probe.free_vram_mib, Some(32_428));
2236 }
2237
2238 #[test]
2239 fn lifecycle_report_plans_llamacpp_unload_and_proves_cleanup() {
2240 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2241 let report = build_lifecycle_report(
2242 profile,
2243 probe_with_used_vram(200),
2244 probe_with_used_vram(820),
2245 Some("http://127.0.0.1:8080/models/unload".to_string()),
2246 Some(42),
2247 DEFAULT_IDLE_TTL_SECONDS,
2248 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
2249 );
2250
2251 assert_eq!(report.lease.mode, LeaseMode::Exclusive);
2252 assert!(report.cleanup.cleanup_proven);
2253 assert_eq!(report.cleanup.status, VramCleanupStatus::Proven);
2254 assert!(report.lease.unload_actions.iter().any(|action| {
2255 action.kind == UnloadActionKind::ProviderApi
2256 && action.endpoint.as_deref() == Some("http://127.0.0.1:8080/models/unload")
2257 }));
2258 assert!(report.lease.unload_actions.iter().any(|action| {
2259 action.kind == UnloadActionKind::ProcessExit
2260 && action.command.as_ref().is_some_and(|command| {
2261 command == &vec!["kill".to_string(), "-TERM".to_string(), "42".to_string()]
2262 })
2263 }));
2264 }
2265
2266 #[test]
2267 fn vram_cleanup_fails_when_provider_process_remains_loaded() {
2268 let pre = probe_with_used_vram(200);
2269 let mut post = probe_with_used_vram(4_000);
2270 post.processes.push(GpuProcess {
2271 pid: Some(42),
2272 process_name: "llama-server".to_string(),
2273 used_memory_mib: Some(3_000),
2274 });
2275
2276 let cleanup = evaluate_vram_cleanup(&pre, &post, DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB);
2277
2278 assert!(!cleanup.cleanup_proven);
2279 assert_eq!(cleanup.status, VramCleanupStatus::NotProven);
2280 assert_eq!(cleanup.blocking_processes.len(), 1);
2281 }
2282
2283 #[test]
2284 fn vram_cleanup_accepts_external_process_accounting() {
2285 let pre = probe_with_used_vram(200);
2286 let mut post = probe_with_used_vram(2_000);
2287 post.processes.push(GpuProcess {
2288 pid: Some(77),
2289 process_name: "python-training-job".to_string(),
2290 used_memory_mib: Some(1_600),
2291 });
2292
2293 let cleanup = evaluate_vram_cleanup(&pre, &post, DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB);
2294
2295 assert!(cleanup.cleanup_proven);
2296 assert_eq!(
2297 cleanup.status,
2298 VramCleanupStatus::ProvenByExternalAccounting
2299 );
2300 assert_eq!(cleanup.external_process_delta_mib, 1_600);
2301 }
2302
2303 #[test]
2304 fn interrupted_run_cleanup_fails_with_orphaned_provider_process() {
2305 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2306 let pre = probe_with_used_vram(200);
2307 let mut post = probe_with_used_vram(8_000);
2308 post.processes.push(GpuProcess {
2309 pid: Some(1234),
2310 process_name: "ollama runner".to_string(),
2311 used_memory_mib: Some(7_000),
2312 });
2313
2314 let report = build_lifecycle_report(
2315 profile,
2316 pre,
2317 post,
2318 None,
2319 Some(1234),
2320 DEFAULT_IDLE_TTL_SECONDS,
2321 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
2322 );
2323
2324 assert!(!report.cleanup.cleanup_proven);
2325 assert_eq!(report.cleanup.status, VramCleanupStatus::NotProven);
2326 assert_eq!(report.cleanup.blocking_processes.len(), 1);
2327 assert!(report.lease.unload_actions.iter().any(|action| {
2328 action.kind == UnloadActionKind::ProcessExit
2329 && action.command.as_ref().is_some_and(|command| {
2330 command == &vec!["kill".to_string(), "-TERM".to_string(), "1234".to_string()]
2331 })
2332 }));
2333 }
2334
2335 fn all_alive(_pid: u32) -> bool {
2338 true
2339 }
2340 fn alive_set(alive: &[u32]) -> impl Fn(u32) -> bool + '_ {
2341 move |pid| alive.contains(&pid)
2342 }
2343
2344 #[test]
2345 fn resolve_lease_file_prefers_explicit_override() {
2346 let path = resolve_lease_file(Some(Path::new("/custom/lease.json")));
2347 assert_eq!(path, PathBuf::from("/custom/lease.json"));
2348 }
2349
2350 #[test]
2351 fn resolve_lease_file_returns_env_value_when_set() {
2352 unsafe {
2356 std::env::set_var(LEASE_FILE_ENV_VAR, "/env/lease.json");
2357 }
2358 let path = resolve_lease_file(None);
2359 unsafe {
2360 std::env::remove_var(LEASE_FILE_ENV_VAR);
2361 }
2362 assert_eq!(path, PathBuf::from("/env/lease.json"));
2363 }
2364
2365 #[test]
2366 fn lease_registry_round_trips_through_json() {
2367 let mut registry = GpuLeaseRegistry::default();
2368 registry.leases.insert(
2369 "qwen3-32b-q4".to_string(),
2370 vec![GpuLeaseRecord {
2371 profile_id: "qwen3-32b-q4".to_string(),
2372 holder_pid: 4242,
2373 holder_command: "tsift".to_string(),
2374 acquired_at_unix_seconds: 100,
2375 lease_mode: LeaseMode::Exclusive,
2376 vram_baseline_mib: 200,
2377 idle_ttl_seconds: 0,
2378 notes: vec!["baseline".to_string()],
2379 }],
2380 );
2381 let payload = serde_json::to_string(®istry).unwrap();
2382 let back: GpuLeaseRegistry = serde_json::from_str(&payload).unwrap();
2383 assert_eq!(registry, back);
2384 assert_eq!(back.version, LEASE_REGISTRY_VERSION);
2385 }
2386
2387 #[test]
2388 fn acquire_exclusive_profile_succeeds_when_free() {
2389 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2390 let mut registry = GpuLeaseRegistry::default();
2391 let acquisition = apply_acquire(
2392 &mut registry,
2393 &profile,
2394 100,
2395 "tsift",
2396 200,
2397 0,
2398 1_000,
2399 all_alive,
2400 );
2401 assert_eq!(acquisition.status, GpuLeaseAcquisitionStatus::Acquired);
2402 assert!(acquisition.conflict.is_none());
2403 assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2404 assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 100);
2405 }
2406
2407 #[test]
2408 fn acquire_exclusive_profile_conflicts_with_live_holder() {
2409 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2410 let mut registry = GpuLeaseRegistry::default();
2411 apply_acquire(
2412 &mut registry,
2413 &profile,
2414 100,
2415 "tsift",
2416 200,
2417 0,
2418 1_000,
2419 all_alive,
2420 );
2421 let second = apply_acquire(
2422 &mut registry,
2423 &profile,
2424 200,
2425 "corky",
2426 250,
2427 0,
2428 1_050,
2429 alive_set(&[100, 200]),
2430 );
2431 assert_eq!(second.status, GpuLeaseAcquisitionStatus::Conflict);
2432 let conflict = second.conflict.unwrap();
2433 assert_eq!(conflict.holder_pid, 100);
2434 assert_eq!(conflict.holder_command, "tsift");
2435 assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2437 assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 100);
2438 }
2439
2440 #[test]
2441 fn acquire_exclusive_profile_reclaims_when_holder_pid_dead() {
2442 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2443 let mut registry = GpuLeaseRegistry::default();
2444 apply_acquire(
2445 &mut registry,
2446 &profile,
2447 100,
2448 "tsift",
2449 200,
2450 0,
2451 1_000,
2452 all_alive,
2453 );
2454 let reclaimed = apply_acquire(
2456 &mut registry,
2457 &profile,
2458 200,
2459 "corky",
2460 250,
2461 0,
2462 1_050,
2463 alive_set(&[200]),
2464 );
2465 assert_eq!(reclaimed.status, GpuLeaseAcquisitionStatus::ReclaimedStale);
2466 assert_eq!(reclaimed.reclaimed.len(), 1);
2467 assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2468 assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 200);
2469 }
2470
2471 #[test]
2472 fn acquire_shared_profile_allows_multiple_live_holders() {
2473 let profile = profile_by_id("qwen3-embedding-0.6b").unwrap();
2474 let mut registry = GpuLeaseRegistry::default();
2475 apply_acquire(
2476 &mut registry,
2477 &profile,
2478 100,
2479 "tsift",
2480 200,
2481 0,
2482 1_000,
2483 all_alive,
2484 );
2485 let second = apply_acquire(
2486 &mut registry,
2487 &profile,
2488 200,
2489 "headroom",
2490 250,
2491 0,
2492 1_050,
2493 alive_set(&[100, 200]),
2494 );
2495 assert_eq!(second.status, GpuLeaseAcquisitionStatus::Acquired);
2496 assert_eq!(registry.leases["qwen3-embedding-0.6b"].len(), 2);
2497 }
2498
2499 #[test]
2500 fn acquire_refreshes_when_same_holder_requests_again() {
2501 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2502 let mut registry = GpuLeaseRegistry::default();
2503 apply_acquire(
2504 &mut registry,
2505 &profile,
2506 100,
2507 "tsift",
2508 200,
2509 0,
2510 1_000,
2511 all_alive,
2512 );
2513 let again = apply_acquire(
2514 &mut registry,
2515 &profile,
2516 100,
2517 "tsift",
2518 180,
2519 0,
2520 1_500,
2521 all_alive,
2522 );
2523 assert_eq!(again.status, GpuLeaseAcquisitionStatus::Refreshed);
2524 assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2525 assert_eq!(
2526 registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
2527 1_500
2528 );
2529 assert_eq!(registry.leases["qwen3-32b-q4"][0].vram_baseline_mib, 180);
2530 }
2531
2532 #[test]
2533 fn acquire_cpu_or_hash_profile_bypasses_registry() {
2534 let profile = profile_by_id("tsift-local-hash-v1").unwrap();
2535 let mut registry = GpuLeaseRegistry::default();
2536 let bypass = apply_acquire(
2537 &mut registry,
2538 &profile,
2539 100,
2540 "tsift",
2541 0,
2542 0,
2543 1_000,
2544 all_alive,
2545 );
2546 assert_eq!(bypass.status, GpuLeaseAcquisitionStatus::CpuOrHashBypass);
2547 assert!(registry.leases.is_empty());
2548 }
2549
2550 #[test]
2551 fn idle_ttl_expires_even_when_pid_still_alive() {
2552 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2553 let mut registry = GpuLeaseRegistry::default();
2554 apply_acquire(
2555 &mut registry,
2556 &profile,
2557 100,
2558 "tsift",
2559 200,
2560 60,
2561 1_000,
2562 all_alive,
2563 );
2564 let reclaimed = apply_acquire(
2566 &mut registry,
2567 &profile,
2568 200,
2569 "corky",
2570 250,
2571 0,
2572 1_120,
2573 all_alive,
2574 );
2575 assert_eq!(reclaimed.status, GpuLeaseAcquisitionStatus::ReclaimedStale);
2576 assert_eq!(registry.leases["qwen3-32b-q4"][0].holder_pid, 200);
2577 }
2578
2579 #[test]
2580 fn release_removes_holder_and_drops_empty_profile() {
2581 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2582 let mut registry = GpuLeaseRegistry::default();
2583 apply_acquire(
2584 &mut registry,
2585 &profile,
2586 100,
2587 "tsift",
2588 200,
2589 0,
2590 1_000,
2591 all_alive,
2592 );
2593 let release = apply_release(&mut registry, "qwen3-32b-q4", 100, 1_050, all_alive);
2594 assert_eq!(release.outcome, GpuLeaseReleaseOutcome::Released);
2595 assert_eq!(release.remaining_holders, 0);
2596 assert!(registry.leases.is_empty());
2597 }
2598
2599 #[test]
2600 fn release_by_non_holder_reports_not_held() {
2601 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2602 let mut registry = GpuLeaseRegistry::default();
2603 apply_acquire(
2604 &mut registry,
2605 &profile,
2606 100,
2607 "tsift",
2608 200,
2609 0,
2610 1_000,
2611 all_alive,
2612 );
2613 let release = apply_release(&mut registry, "qwen3-32b-q4", 999, 1_050, all_alive);
2614 assert_eq!(release.outcome, GpuLeaseReleaseOutcome::NotHeld);
2615 assert_eq!(registry.leases["qwen3-32b-q4"].len(), 1);
2616 }
2617
2618 #[test]
2619 fn acquire_and_release_round_trip_through_file() {
2620 let dir = tempfile_dir();
2621 let path = dir.join("gpu-lease.json");
2622 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2623
2624 let mut registry = GpuLeaseRegistry::default();
2625 let acquisition = apply_acquire(
2626 &mut registry,
2627 &profile,
2628 4242,
2629 "tsift",
2630 220,
2631 0,
2632 1_000,
2633 all_alive,
2634 );
2635 assert_eq!(acquisition.status, GpuLeaseAcquisitionStatus::Acquired);
2636 write_lease_registry(&path, ®istry).unwrap();
2637
2638 let read_back = read_lease_registry(&path).unwrap();
2639 assert_eq!(read_back, registry);
2640 assert_eq!(read_back.leases["qwen3-32b-q4"][0].holder_pid, 4242);
2641
2642 let release = apply_release(&mut registry, "qwen3-32b-q4", 4242, 1_050, all_alive);
2643 assert_eq!(release.outcome, GpuLeaseReleaseOutcome::Released);
2644 write_lease_registry(&path, ®istry).unwrap();
2645
2646 let after = read_lease_registry(&path).unwrap();
2647 assert!(after.leases.is_empty());
2648 }
2649
2650 #[test]
2651 fn read_lease_registry_returns_default_for_missing_file() {
2652 let path = Path::new("/definitely/not/a/real/path/lease.json");
2653 let registry = read_lease_registry(path).unwrap();
2654 assert_eq!(registry, GpuLeaseRegistry::default());
2655 }
2656
2657 #[test]
2658 fn registry_lock_path_appends_lock_suffix() {
2659 assert_eq!(
2660 registry_lock_path(Path::new("/tmp/x/gpu-lease.json")),
2661 PathBuf::from("/tmp/x/gpu-lease.json.lock")
2662 );
2663 }
2664
2665 #[test]
2666 fn acquire_lease_creates_sidecar_lock_file() {
2667 let dir = tempfile_dir();
2668 let path = dir.join("gpu-lease.json");
2669 acquire_lease("qwen3-32b-q4", std::process::id(), "tsift", 0, 0, 1_000, &path).unwrap();
2672 assert!(
2673 registry_lock_path(&path).exists(),
2674 "sidecar lock file should exist after a locked acquire"
2675 );
2676 }
2677
2678 #[test]
2679 fn apply_renew_slides_heartbeat_so_ttl_holder_survives() {
2680 let mut registry = GpuLeaseRegistry::default();
2681 let profile = profile_by_id("qwen3-32b-q4").unwrap();
2682 apply_acquire(&mut registry, &profile, 100, "tsift", 200, 100, 1_000, all_alive);
2684 let renew = apply_renew(&mut registry, "qwen3-32b-q4", 100, 1_050, all_alive);
2686 assert_eq!(renew.outcome, GpuLeaseRenewOutcome::Renewed);
2687 assert_eq!(renew.renewed_at_unix_seconds, Some(1_050));
2688 assert_eq!(
2689 registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
2690 1_050
2691 );
2692 let pruned = prune_stale_leases(&mut registry, 1_120, all_alive);
2696 assert!(pruned.is_empty());
2697 assert!(registry.leases.contains_key("qwen3-32b-q4"));
2698 }
2699
2700 #[test]
2701 fn apply_renew_reports_profile_absent_for_unheld_profile() {
2702 let mut registry = GpuLeaseRegistry::default();
2703 let renew = apply_renew(&mut registry, "qwen3-32b-q4", 100, 1_000, all_alive);
2704 assert_eq!(renew.outcome, GpuLeaseRenewOutcome::ProfileAbsent);
2705 assert!(renew.renewed_at_unix_seconds.is_none());
2706 }
2707
2708 #[test]
2709 fn reap_leases_reclaims_dead_pid_and_reports_emptied_profile() {
2710 let dir = tempfile_dir();
2711 let path = dir.join("gpu-lease.json");
2712 let mut registry = GpuLeaseRegistry::default();
2713 registry.leases.insert(
2714 "qwen3-32b-q4".to_string(),
2715 vec![GpuLeaseRecord {
2716 profile_id: "qwen3-32b-q4".to_string(),
2717 holder_pid: 4_000_000_000,
2720 holder_command: "crashed-session".to_string(),
2721 acquired_at_unix_seconds: 1_000,
2722 lease_mode: LeaseMode::Exclusive,
2723 vram_baseline_mib: 200,
2724 idle_ttl_seconds: 0,
2725 notes: Vec::new(),
2726 }],
2727 );
2728 write_lease_registry(&path, ®istry).unwrap();
2729
2730 let reap = reap_leases(2_000, &path).unwrap();
2731 assert_eq!(reap.reclaimed.len(), 1);
2732 assert_eq!(reap.emptied_profiles, vec!["qwen3-32b-q4".to_string()]);
2733 let after = read_lease_registry(&path).unwrap();
2734 assert!(after.leases.is_empty());
2735 }
2736
2737 #[test]
2738 fn renew_lease_round_trips_through_file() {
2739 let dir = tempfile_dir();
2740 let path = dir.join("gpu-lease.json");
2741 let pid = std::process::id();
2742 acquire_lease("qwen3-32b-q4", pid, "tsift", 0, 0, 1_000, &path).unwrap();
2745 let renew = renew_lease("qwen3-32b-q4", pid, 5_000, &path).unwrap();
2746 assert_eq!(renew.outcome, GpuLeaseRenewOutcome::Renewed);
2747 let registry = read_lease_registry(&path).unwrap();
2748 assert_eq!(
2749 registry.leases["qwen3-32b-q4"][0].acquired_at_unix_seconds,
2750 5_000
2751 );
2752 }
2753
2754 #[test]
2755 fn prune_stale_leaves_healthy_entries_alone() {
2756 let mut registry = GpuLeaseRegistry::default();
2757 registry.leases.insert(
2758 "qwen3-32b-q4".to_string(),
2759 vec![GpuLeaseRecord {
2760 profile_id: "qwen3-32b-q4".to_string(),
2761 holder_pid: 100,
2762 holder_command: "tsift".to_string(),
2763 acquired_at_unix_seconds: 1_000,
2764 lease_mode: LeaseMode::Exclusive,
2765 vram_baseline_mib: 200,
2766 idle_ttl_seconds: 0,
2767 notes: Vec::new(),
2768 }],
2769 );
2770 let pruned = prune_stale_leases(&mut registry, 1_010, alive_set(&[100]));
2771 assert!(pruned.is_empty());
2772 assert!(registry.leases.contains_key("qwen3-32b-q4"));
2773 }
2774
2775 fn env_lock() -> std::sync::MutexGuard<'static, ()> {
2778 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2779 ENV_LOCK
2780 .lock()
2781 .unwrap_or_else(|poisoned| poisoned.into_inner())
2782 }
2783
2784 fn tempfile_dir() -> PathBuf {
2785 use std::sync::atomic::{AtomicU64, Ordering};
2786 static COUNTER: AtomicU64 = AtomicU64::new(0);
2787 let dir = std::env::temp_dir().join(format!(
2790 "tsift-lease-test-{}-{}-{}",
2791 std::process::id(),
2792 current_unix_seconds(),
2793 COUNTER.fetch_add(1, Ordering::Relaxed)
2794 ));
2795 std::fs::create_dir_all(&dir).unwrap();
2796 dir
2797 }
2798
2799 #[test]
2802 fn profile_preference_parses_cli_value() {
2803 assert_eq!(ProfilePreference::from_cli(None), ProfilePreference::Auto);
2804 assert_eq!(
2805 ProfilePreference::from_cli(Some("")),
2806 ProfilePreference::Auto
2807 );
2808 assert_eq!(
2809 ProfilePreference::from_cli(Some("hash")),
2810 ProfilePreference::ForceHash
2811 );
2812 assert_eq!(
2813 ProfilePreference::from_cli(Some("tsift-local-hash-v1")),
2814 ProfilePreference::ForceHash
2815 );
2816 assert_eq!(
2817 ProfilePreference::from_cli(Some("qwen3-32b-q4")),
2818 ProfilePreference::Pinned("qwen3-32b-q4".to_string())
2819 );
2820 }
2821
2822 #[test]
2823 fn resolve_auto_picks_recommended_gpu_profile_on_clear_5090() {
2824 let probe = rtx_5090_probe();
2825 let resolution =
2826 resolve_profile_preference(&ProfilePreference::Auto, ModelRole::Extract, &probe);
2827 assert_eq!(resolution.source, ProfileResolutionSource::AutoRanked);
2828 assert!(resolution.selectable);
2829 assert_eq!(resolution.profile.id, "qwen3-32b-q4");
2830 }
2831
2832 #[test]
2833 fn resolve_auto_falls_back_to_hash_when_gpu_unavailable() {
2834 let probe = GpuProbe::unavailable("missing");
2835 let resolution =
2836 resolve_profile_preference(&ProfilePreference::Auto, ModelRole::Extract, &probe);
2837 assert_eq!(resolution.source, ProfileResolutionSource::AutoRanked);
2838 assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2839 assert!(resolution.selectable);
2840 assert!(resolution.reason.contains("no GPU profile selectable"));
2841 }
2842
2843 #[test]
2844 fn resolve_pinned_selectable_profile_is_used_as_is() {
2845 let probe = rtx_5090_probe();
2846 let resolution = resolve_profile_preference(
2847 &ProfilePreference::Pinned("qwen3-embedding-0.6b".to_string()),
2848 ModelRole::Embed,
2849 &probe,
2850 );
2851 assert_eq!(resolution.source, ProfileResolutionSource::Pinned);
2852 assert_eq!(resolution.profile.id, "qwen3-embedding-0.6b");
2853 assert!(resolution.selectable);
2854 }
2855
2856 #[test]
2857 fn resolve_pinned_profile_with_wrong_role_falls_back_to_hash() {
2858 let probe = rtx_5090_probe();
2859 let resolution = resolve_profile_preference(
2860 &ProfilePreference::Pinned("qwen3-embedding-0.6b".to_string()),
2861 ModelRole::Extract,
2862 &probe,
2863 );
2864 assert_eq!(
2865 resolution.source,
2866 ProfileResolutionSource::PinnedUnselectable
2867 );
2868 assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2869 assert!(resolution.reason.contains("does not support role"));
2870 }
2871
2872 #[test]
2873 fn resolve_pinned_profile_that_does_not_fit_vram_falls_back_to_hash() {
2874 let probe = probe_with_used_vram(30_000);
2876 let resolution = resolve_profile_preference(
2877 &ProfilePreference::Pinned("qwen3-32b-q4".to_string()),
2878 ModelRole::Extract,
2879 &probe,
2880 );
2881 assert_eq!(
2882 resolution.source,
2883 ProfileResolutionSource::PinnedUnselectable
2884 );
2885 assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2886 assert!(resolution.selectable);
2887 assert!(resolution.reason.contains("not selectable"));
2888 }
2889
2890 #[test]
2891 fn resolve_force_hash_always_uses_hash_profile() {
2892 let probe = rtx_5090_probe();
2893 let resolution =
2894 resolve_profile_preference(&ProfilePreference::ForceHash, ModelRole::Extract, &probe);
2895 assert_eq!(resolution.source, ProfileResolutionSource::ForcedHash);
2896 assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2897 assert!(resolution.selectable);
2898 assert!(resolution.reason.contains("forced"));
2899 }
2900
2901 #[test]
2902 fn resolve_pinned_unknown_profile_id_falls_back_to_hash() {
2903 let probe = rtx_5090_probe();
2904 let resolution = resolve_profile_preference(
2905 &ProfilePreference::Pinned("not-a-real-profile".to_string()),
2906 ModelRole::Embed,
2907 &probe,
2908 );
2909 assert_eq!(
2910 resolution.source,
2911 ProfileResolutionSource::PinnedUnselectable
2912 );
2913 assert_eq!(resolution.profile.id, "tsift-local-hash-v1");
2914 assert!(resolution.reason.contains("unknown"));
2915 }
2916
2917 #[test]
2920 fn resolve_endpoint_returns_explicit_override_for_any_strategy() {
2921 for strategy in [
2922 UnloadStrategy::LlamaCppRouterUnload,
2923 UnloadStrategy::OllamaKeepAliveZero,
2924 UnloadStrategy::VllmSleep,
2925 UnloadStrategy::ProcessExit,
2926 UnloadStrategy::None,
2927 ] {
2928 let resolved = resolve_provider_endpoint(&strategy, Some("http://custom:9999/path"));
2929 assert_eq!(
2930 resolved, "http://custom:9999/path",
2931 "explicit override should win for {strategy:?}"
2932 );
2933 }
2934 }
2935
2936 #[test]
2937 fn resolve_endpoint_uses_compile_time_default_when_no_env_no_explicit() {
2938 let _env = env_lock();
2940 unsafe {
2943 std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
2944 std::env::remove_var(OLLAMA_ENDPOINT_ENV_VAR);
2945 std::env::remove_var(VLLM_ENDPOINT_ENV_VAR);
2946 }
2947 assert_eq!(
2948 resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None),
2949 DEFAULT_LLAMA_CPP_ENDPOINT
2950 );
2951 assert_eq!(
2952 resolve_provider_endpoint(&UnloadStrategy::OllamaKeepAliveZero, None),
2953 DEFAULT_OLLAMA_ENDPOINT
2954 );
2955 assert_eq!(
2956 resolve_provider_endpoint(&UnloadStrategy::VllmSleep, None),
2957 DEFAULT_VLLM_ENDPOINT
2958 );
2959 assert_eq!(
2960 resolve_provider_endpoint(&UnloadStrategy::ProcessExit, None),
2961 ""
2962 );
2963 assert_eq!(resolve_provider_endpoint(&UnloadStrategy::None, None), "");
2964 }
2965
2966 #[test]
2967 fn resolve_endpoint_env_var_overrides_default_for_llama_cpp() {
2968 let _env = env_lock();
2969 unsafe {
2971 std::env::set_var(
2972 LLAMA_CPP_ENDPOINT_ENV_VAR,
2973 "http://127.0.0.1:8081/models/unload",
2974 );
2975 }
2976 let resolved = resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None);
2977 unsafe {
2978 std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
2979 }
2980 assert_eq!(resolved, "http://127.0.0.1:8081/models/unload");
2981 }
2982
2983 #[test]
2984 fn resolve_endpoint_blank_env_var_falls_back_to_default() {
2985 let _env = env_lock();
2986 unsafe {
2988 std::env::set_var(LLAMA_CPP_ENDPOINT_ENV_VAR, " ");
2989 }
2990 let resolved = resolve_provider_endpoint(&UnloadStrategy::LlamaCppRouterUnload, None);
2991 unsafe {
2992 std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
2993 }
2994 assert_eq!(resolved, DEFAULT_LLAMA_CPP_ENDPOINT);
2995 }
2996
2997 #[test]
2998 fn build_unload_actions_picks_up_env_var_for_llama_cpp_endpoint() {
2999 let _env = env_lock();
3000 let profile = profile_by_id("qwen3-32b-q4").unwrap();
3001 unsafe {
3003 std::env::set_var(
3004 LLAMA_CPP_ENDPOINT_ENV_VAR,
3005 "http://127.0.0.1:8081/models/unload",
3006 );
3007 }
3008 let actions = build_unload_actions(&profile, None, Some(42));
3009 unsafe {
3010 std::env::remove_var(LLAMA_CPP_ENDPOINT_ENV_VAR);
3011 }
3012 let unload_action = actions
3013 .iter()
3014 .find(|action| action.kind == UnloadActionKind::ProviderApi)
3015 .expect("provider api action present");
3016 assert_eq!(
3017 unload_action.endpoint.as_deref(),
3018 Some("http://127.0.0.1:8081/models/unload")
3019 );
3020 }
3021
3022 fn probe_pair(pre_used: u64, post_used: u64) -> (GpuProbe, GpuProbe) {
3025 (
3026 probe_with_used_vram(pre_used),
3027 probe_with_used_vram(post_used),
3028 )
3029 }
3030
3031 #[test]
3032 fn swap_to_same_profile_is_noop() {
3033 let from = profile_by_id("qwen3-32b-q4").unwrap();
3034 let to = profile_by_id("qwen3-32b-q4").unwrap();
3035 let (pre, post) = probe_pair(200, 200);
3036 let report = build_swap_report(
3037 from,
3038 to,
3039 pre,
3040 post,
3041 None,
3042 None,
3043 DEFAULT_IDLE_TTL_SECONDS,
3044 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3045 );
3046 assert_eq!(report.swap_status, SwapStatus::NoOpSameProfile);
3047 }
3048
3049 #[test]
3050 fn swap_from_big_to_small_embedding_when_cleanup_proven_is_swapped() {
3051 let from = profile_by_id("qwen3-32b-q4").unwrap();
3052 let to = profile_by_id("qwen3-embedding-0.6b").unwrap();
3053 let (pre, post) = probe_pair(28_000, 200);
3055 let report = build_swap_report(
3056 from,
3057 to,
3058 pre,
3059 post,
3060 None,
3061 Some(42),
3062 DEFAULT_IDLE_TTL_SECONDS,
3063 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3064 );
3065 assert_eq!(report.swap_status, SwapStatus::Swapped);
3066 assert!(report.unload.cleanup.cleanup_proven);
3067 assert_eq!(report.target_resolution.profile.id, "qwen3-embedding-0.6b");
3068 }
3069
3070 #[test]
3071 fn swap_to_hash_fallback_is_swapped_to_hash_when_cleanup_proven() {
3072 let from = profile_by_id("qwen3-32b-q4").unwrap();
3073 let to = profile_by_id("tsift-local-hash-v1").unwrap();
3074 let (pre, post) = probe_pair(28_000, 200);
3075 let report = build_swap_report(
3076 from,
3077 to,
3078 pre,
3079 post,
3080 None,
3081 Some(42),
3082 DEFAULT_IDLE_TTL_SECONDS,
3083 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3084 );
3085 assert_eq!(report.swap_status, SwapStatus::SwappedToHash);
3086 assert!(report.unload.cleanup.cleanup_proven);
3087 }
3088
3089 #[test]
3090 fn swap_blocks_when_source_unload_not_proven() {
3091 let from = profile_by_id("qwen3-32b-q4").unwrap();
3092 let to = profile_by_id("qwen3-embedding-0.6b").unwrap();
3093 let pre = probe_with_used_vram(200);
3096 let mut post = probe_with_used_vram(8_000);
3097 post.processes.push(GpuProcess {
3098 pid: Some(42),
3099 process_name: "llama-server".to_string(),
3100 used_memory_mib: Some(7_000),
3101 });
3102 let report = build_swap_report(
3103 from,
3104 to,
3105 pre,
3106 post,
3107 None,
3108 Some(42),
3109 DEFAULT_IDLE_TTL_SECONDS,
3110 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3111 );
3112 assert_eq!(report.swap_status, SwapStatus::UnloadNotProven);
3113 assert!(!report.unload.cleanup.cleanup_proven);
3114 assert!(
3115 report
3116 .notes
3117 .iter()
3118 .any(|note| note.contains("DO NOT load target"))
3119 );
3120 }
3121
3122 #[test]
3123 fn swap_reports_target_unselectable_when_post_unload_vram_still_high() {
3124 let from = profile_by_id("qwen3-32b-q4").unwrap();
3125 let to = profile_by_id("qwen3-32b-q4").unwrap();
3126 let pre = probe_with_used_vram(29_500);
3130 let post = probe_with_used_vram(29_600);
3131 let report = build_swap_report(
3132 from,
3133 to,
3134 pre,
3135 post,
3136 None,
3137 None,
3138 DEFAULT_IDLE_TTL_SECONDS,
3139 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3140 );
3141 assert_eq!(report.swap_status, SwapStatus::NoOpSameProfile);
3143 let from = profile_by_id("qwen3-32b-q4").unwrap();
3145 let to = profile_by_id("qwen3-embedding-8b").unwrap();
3146 let pre = probe_with_used_vram(30_000);
3147 let post = probe_with_used_vram(30_500);
3148 let report = build_swap_report(
3149 from,
3150 to,
3151 pre,
3152 post,
3153 None,
3154 None,
3155 DEFAULT_IDLE_TTL_SECONDS,
3156 DEFAULT_VRAM_CLEANUP_TOLERANCE_MIB,
3157 );
3158 assert_eq!(
3159 report.swap_status,
3160 SwapStatus::UnloadProvenTargetUnselectable
3161 );
3162 assert!(
3163 report
3164 .notes
3165 .iter()
3166 .any(|note| note.contains("not selectable on the post-unload probe"))
3167 );
3168 }
3169
3170 #[test]
3175 fn rewrite_unload_body_model_replaces_model_field() {
3176 let original = r#"{"model":"qwen3-32b-q4-ollama-default","prompt":"","keep_alive":0}"#;
3177 let rewritten = rewrite_unload_body_model(original, "hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M");
3178 let value: serde_json::Value =
3179 serde_json::from_str(&rewritten).expect("rewritten body is valid JSON");
3180 assert_eq!(
3181 value["model"].as_str(),
3182 Some("hf.co/Qwen/Qwen3-32B-GGUF:Q4_K_M")
3183 );
3184 assert_eq!(value["keep_alive"].as_i64(), Some(0));
3186 assert_eq!(value["prompt"].as_str(), Some(""));
3187 }
3188
3189 #[test]
3190 fn rewrite_unload_body_model_preserves_body_when_override_is_empty() {
3191 let original = r#"{"model":"default-tag","keep_alive":0}"#;
3192 let rewritten = rewrite_unload_body_model(original, "");
3193 assert_eq!(rewritten, original);
3195 }
3196
3197 #[test]
3198 fn rewrite_unload_body_model_falls_back_on_invalid_json() {
3199 let original = "not valid json {{{";
3200 let rewritten = rewrite_unload_body_model(original, "any-tag");
3201 assert_eq!(rewritten, original);
3202 }
3203
3204 #[test]
3205 fn normalize_unload_url_appends_generate_path_for_bare_host() {
3206 let url = normalize_unload_url("http://127.0.0.1:11434");
3207 assert_eq!(url, "http://127.0.0.1:11434/api/generate");
3208 }
3209
3210 #[test]
3211 fn normalize_unload_url_idempotent_for_full_generate_url() {
3212 let url = normalize_unload_url("http://127.0.0.1:11434/api/generate");
3213 assert_eq!(url, "http://127.0.0.1:11434/api/generate");
3214 }
3215
3216 #[test]
3217 fn normalize_unload_url_strips_trailing_slash() {
3218 let url = normalize_unload_url("http://127.0.0.1:11434/");
3219 assert_eq!(url, "http://127.0.0.1:11434/api/generate");
3220 }
3221
3222 #[test]
3223 fn prepare_unload_request_returns_none_for_non_api_actions() {
3224 let noop = ProviderUnloadAction {
3225 kind: UnloadActionKind::Noop,
3226 label: "noop".to_string(),
3227 command: None,
3228 http_method: None,
3229 endpoint: None,
3230 body_json: None,
3231 required: false,
3232 };
3233 assert!(prepare_unload_request(&noop, "any-tag").is_none());
3234 }
3235
3236 #[test]
3237 fn prepare_unload_request_applies_model_override_to_body() {
3238 let action = ProviderUnloadAction {
3242 kind: UnloadActionKind::ProviderApi,
3243 label: "ollama keep_alive zero".to_string(),
3244 command: Some(vec!["ollama".to_string(), "stop".to_string()]),
3245 http_method: Some("POST".to_string()),
3246 endpoint: Some("http://127.0.0.1:11434".to_string()),
3247 body_json: Some(
3248 r#"{"model":"profile-default-tag","prompt":"","keep_alive":0}"#.to_string(),
3249 ),
3250 required: true,
3251 };
3252 let req = prepare_unload_request(&action, "override-tag")
3253 .expect("ProviderApi action prepares a request");
3254 assert_eq!(req.url, "http://127.0.0.1:11434/api/generate");
3255 assert!(req.body.contains("\"model\":\"override-tag\""));
3256 assert!(!req.body.contains("profile-default-tag"));
3257 assert_eq!(
3258 req.fallback_command,
3259 Some(vec!["ollama".to_string(), "stop".to_string()])
3260 );
3261 }
3262
3263 #[test]
3264 fn prepare_unload_request_synthesizes_body_when_plan_has_none() {
3265 let action = ProviderUnloadAction {
3266 kind: UnloadActionKind::ProviderApi,
3267 label: "synthesized".to_string(),
3268 command: None,
3269 http_method: Some("POST".to_string()),
3270 endpoint: Some("http://host:11434".to_string()),
3271 body_json: None,
3272 required: true,
3273 };
3274 let req = prepare_unload_request(&action, "synth-tag").unwrap();
3275 assert!(req.body.contains("\"model\":\"synth-tag\""));
3276 assert!(req.body.contains("\"keep_alive\":0"));
3277 }
3278
3279 #[test]
3280 fn execute_unload_actions_reports_non_api_as_skipped() {
3281 let actions = vec![ProviderUnloadAction {
3282 kind: UnloadActionKind::Noop,
3283 label: "no GPU unload required".to_string(),
3284 command: None,
3285 http_method: None,
3286 endpoint: None,
3287 body_json: None,
3288 required: false,
3289 }];
3290 let results = execute_unload_actions(&actions, "any-tag");
3291 assert_eq!(results.len(), 1);
3292 assert!(!results[0].executed);
3293 assert_eq!(results[0].outcome, "skipped: non-API action");
3294 }
3295}