diff --git a/crates/core/src/tuning.rs b/crates/core/src/tuning.rs index 7bf9029..c3a2417 100644 --- a/crates/core/src/tuning.rs +++ b/crates/core/src/tuning.rs @@ -3,6 +3,8 @@ //! callable from both inference call sites. use crate::power::{self, PowerState}; +use std::collections::HashSet; +use std::sync::{Mutex, OnceLock}; /// Lower bound for inference threads. Single-threaded inference is /// measurably worse than two on every multi-core part. @@ -17,7 +19,7 @@ pub const MIN_INFERENCE_THREADS: usize = 2; /// MAGNOTIA_INFERENCE_THREADS. pub const MAX_INFERENCE_THREADS: usize = 8; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum Workload { /// Llama-style transformer. Near-zero CPU work when fully /// offloaded; CPU floor in that case is GPU_FLOOR_LLM. @@ -31,6 +33,14 @@ pub enum Workload { const GPU_FLOOR_LLM: usize = 2; const GPU_FLOOR_WHISPER: usize = 4; +/// Whether a given (workload, on_battery, gpu_offloaded) tuple has +/// already been logged this process. Prevents log spam when the same +/// configuration is exercised on every inference call. +fn log_seen() -> &'static Mutex> { + static SEEN: OnceLock>> = OnceLock::new(); + SEEN.get_or_init(|| Mutex::new(HashSet::new())) +} + /// Inference thread count, clamped to the physical-core budget plus /// the battery and GPU-offload heuristics. /// @@ -56,17 +66,41 @@ pub fn inference_thread_count(workload: Workload, gpu_offloaded: bool) -> usize .map(|p| p.get()) .unwrap_or(MIN_INFERENCE_THREADS) }; - if power::probe_power_state() == PowerState::OnBattery { + let on_battery = power::probe_power_state() == PowerState::OnBattery; + let mut clamps: Vec<&'static str> = Vec::new(); + if on_battery { chosen /= 2; + clamps.push("battery"); } if gpu_offloaded { let floor = match workload { Workload::Llm => GPU_FLOOR_LLM, Workload::Whisper => GPU_FLOOR_WHISPER, }; - chosen = chosen.min(floor); + if chosen > floor { + chosen = floor; + clamps.push("gpu"); + } } - chosen.clamp(MIN_INFERENCE_THREADS, MAX_INFERENCE_THREADS) + let final_value = chosen.clamp(MIN_INFERENCE_THREADS, MAX_INFERENCE_THREADS); + + // Log once per (workload, on_battery, gpu_offloaded) tuple. + let key = (workload, on_battery, gpu_offloaded); + if let Ok(mut seen) = log_seen().lock() { + if seen.insert(key) { + tracing::info!( + target: "magnotia_core::tuning", + threads = final_value, + workload = ?workload, + gpu_offloaded, + on_battery, + clamps = ?clamps, + "inference_thread_count" + ); + } + } + + final_value } #[cfg(test)] @@ -182,4 +216,18 @@ mod tests { }); }); } + + #[test] + fn logging_does_not_panic() { + // Smoke: helper should run without panicking when logging is + // wired. This is covered by the other tests too, but kept + // explicitly to document the behaviour. + with_thread_env_lock(|| { + std::env::remove_var("MAGNOTIA_INFERENCE_THREADS"); + crate::power::with_override(Some(PowerState::OnBattery), || { + let _ = inference_thread_count(Workload::Llm, true); + let _ = inference_thread_count(Workload::Whisper, false); + }); + }); + } }