Lumotia/crates/core/src/tuning.rs

//! Inference thread-count tuning. Combines physical-core budget,
//! battery awareness, and GPU-offload awareness into a single helper
//! 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.
pub const MIN_INFERENCE_THREADS: usize = 2;

/// Upper bound for inference threads. whisper.cpp + llama.cpp scaling
/// flattens around physical core count; SMT siblings contend for FPU
/// during F16/F32 matmul, so going past physical cores anti-scales.
/// 8 is a conservative ceiling that leaves <5% on the table for
/// big-iron desktops while keeping consumer 6c/12t laptops out of
/// contention territory. Users can override at runtime via
/// MAGNOTIA_INFERENCE_THREADS.
pub const MAX_INFERENCE_THREADS: usize = 8;

#[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.
    Llm,
    /// Whisper transcription. Keeps mel spectrogram, decoder
    /// bookkeeping, and beam search on the CPU even with full Vulkan
    /// offload; CPU floor is higher: GPU_FLOOR_WHISPER.
    Whisper,
}

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<HashSet<(Workload, bool, bool)>> {
    static SEEN: OnceLock<Mutex<HashSet<(Workload, bool, bool)>>> = 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.
///
/// Resolution order:
/// 1. `MAGNOTIA_INFERENCE_THREADS=N` — absolute bypass, returns N.
/// 2. base = num_cpus::get_physical() (fallback: available_parallelism).
/// 3. on battery → base /= 2.
/// 4. gpu_offloaded → base = min(base, gpu_floor(workload)).
/// 5. clamp to [MIN_INFERENCE_THREADS, MAX_INFERENCE_THREADS].
pub fn inference_thread_count(workload: Workload, gpu_offloaded: bool) -> usize {
    if let Ok(s) = std::env::var("MAGNOTIA_INFERENCE_THREADS") {
        if let Ok(n) = s.parse::<usize>() {
            if n > 0 {
                return n;
            }
        }
    }
    let physical = num_cpus::get_physical();
    let mut chosen = if physical > 0 {
        physical
    } else {
        std::thread::available_parallelism()
            .map(|p| p.get())
            .unwrap_or(MIN_INFERENCE_THREADS)
    };
    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,
        };
        if chosen > floor {
            chosen = floor;
            clamps.push("gpu");
        }
    }
    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)]
mod tests {
    use super::*;

    /// Serialises tests that read/write `MAGNOTIA_INFERENCE_THREADS` so
    /// they don't race under cargo's parallel test runner.
    /// Mirrors the pattern used by `power::with_override`.
    fn with_thread_env_lock<R>(body: impl FnOnce() -> R) -> R {
        use std::sync::Mutex;
        static THREAD_ENV_LOCK: Mutex<()> = Mutex::new(());
        let _lock = THREAD_ENV_LOCK.lock().expect("tuning thread-env lock poisoned");
        body()
    }

    #[test]
    fn matches_existing_clamp_when_no_clamps_apply() {
        // With no env override, no battery, no gpu_offload, helper
        // should return physical-core count clamped to [2, 8].
        // We can't pin physical exactly without mocking num_cpus; just
        // assert the result is in range.
        with_thread_env_lock(|| {
            std::env::remove_var("MAGNOTIA_INFERENCE_THREADS");
            let n = inference_thread_count(Workload::Llm, false);
            assert!(n >= MIN_INFERENCE_THREADS && n <= MAX_INFERENCE_THREADS,
                    "got {n}, expected within [{MIN_INFERENCE_THREADS}, {MAX_INFERENCE_THREADS}]");
        });
    }

    #[test]
    fn env_var_bypasses_clamps() {
        with_thread_env_lock(|| {
            std::env::set_var("MAGNOTIA_INFERENCE_THREADS", "10");
            let n = inference_thread_count(Workload::Llm, true);
            assert_eq!(n, 10);
            std::env::remove_var("MAGNOTIA_INFERENCE_THREADS");
        });
    }

    #[test]
    fn workload_variants_distinct() {
        assert_ne!(Workload::Llm, Workload::Whisper);
    }

    #[test]
    fn battery_halves_thread_count() {
        with_thread_env_lock(|| {
            std::env::remove_var("MAGNOTIA_INFERENCE_THREADS");
            // Measure on battery, then on AC — sequential, not nested,
            // to avoid re-entrant deadlock on power::TEST_LOCK.
            let on_battery = crate::power::with_override(Some(PowerState::OnBattery), || {
                inference_thread_count(Workload::Llm, false)
            });
            let on_ac = crate::power::with_override(Some(PowerState::OnAc), || {
                inference_thread_count(Workload::Llm, false)
            });
            // on_battery should be roughly half of on_ac, clamped to MIN.
            if on_ac > MIN_INFERENCE_THREADS {
                assert!(
                    on_battery <= on_ac,
                    "battery ({on_battery}) should be <= AC ({on_ac})"
                );
                assert!(on_battery >= MIN_INFERENCE_THREADS);
            }
        });
    }

    #[test]
    fn gpu_offload_clamps_llm_to_floor() {
        with_thread_env_lock(|| {
            std::env::remove_var("MAGNOTIA_INFERENCE_THREADS");
            crate::power::with_override(Some(PowerState::OnAc), || {
                let n = inference_thread_count(Workload::Llm, true);
                assert!(n <= GPU_FLOOR_LLM.max(MIN_INFERENCE_THREADS),
                        "Llm with GPU offload should clamp to {GPU_FLOOR_LLM}, got {n}");
                assert!(n >= MIN_INFERENCE_THREADS);
            });
        });
    }

    #[test]
    fn gpu_offload_clamps_whisper_to_floor() {
        with_thread_env_lock(|| {
            std::env::remove_var("MAGNOTIA_INFERENCE_THREADS");
            crate::power::with_override(Some(PowerState::OnAc), || {
                let n = inference_thread_count(Workload::Whisper, true);
                // Whisper floor is 4 on machines with >=4 physical cores;
                // can't go lower than physical so on a 2c host we stay at 2.
                assert!(n <= GPU_FLOOR_WHISPER, "got {n}");
                assert!(n >= MIN_INFERENCE_THREADS);
            });
        });
    }

    #[test]
    fn whisper_gpu_floor_is_at_least_llm_gpu_floor() {
        // Architectural invariant: whisper retains CPU work even when
        // GPU-offloaded; its floor must not be lower than the LLM's.
        assert!(GPU_FLOOR_WHISPER >= GPU_FLOOR_LLM);
    }

    #[test]
    fn gpu_offload_off_does_not_clamp_below_battery_calc() {
        with_thread_env_lock(|| {
            std::env::remove_var("MAGNOTIA_INFERENCE_THREADS");
            // Sequential measurements; with_override is non-reentrant.
            crate::power::with_override(Some(PowerState::OnAc), || {
                let no_gpu = inference_thread_count(Workload::Llm, false);
                let with_gpu = inference_thread_count(Workload::Llm, true);
                // GPU-offloaded should be <= non-offloaded.
                assert!(with_gpu <= no_gpu);
            });
        });
    }

    #[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);
            });
        });
    }
}