Lumotia/crates/transcription/tests/thread_sweep.rs

//! Thread-count scaling sweep for Whisper Tiny.
//! Runs the JFK clip at n_threads = 1, 2, 4, 6, 8, 12, prints RTF tables.
//! Env-gated by `MAGNOTIA_WHISPER_TEST_MODEL` + `MAGNOTIA_WHISPER_TEST_AUDIO`.
//!
//! Now prints multiple panels driven by `MAGNOTIA_POWER_STATE_OVERRIDE` so
//! the helper's predicted thread count for each (power, GPU) combination
//! can be compared against the empirical RTF data.

use std::env;
use std::time::Instant;

use magnotia_core::hardware::vulkan_loader_available;
use magnotia_core::tuning::{inference_thread_count, Workload};
use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperContextParameters};

#[test]
fn whisper_thread_count_sweep() {
    let Ok(model_path) = env::var("MAGNOTIA_WHISPER_TEST_MODEL") else {
        return;
    };
    let Ok(audio_path) = env::var("MAGNOTIA_WHISPER_TEST_AUDIO") else {
        return;
    };

    let bytes = std::fs::read(&audio_path).expect("read wav");
    let sample_rate = u32::from_le_bytes(bytes[24..28].try_into().unwrap());
    let pcm = &bytes[44..];
    let samples: Vec<f32> = pcm
        .chunks_exact(2)
        .map(|c| i16::from_le_bytes([c[0], c[1]]) as f32 / 32768.0)
        .collect();
    let audio_secs = samples.len() as f64 / sample_rate as f64;
    eprintln!("[sweep] audio: {:.2}s @ {} Hz", audio_secs, sample_rate);

    let logical = num_cpus::get();
    let physical = num_cpus::get_physical();
    eprintln!("[sweep] CPU: physical={}, logical={}", physical, logical);

    let ctx = WhisperContext::new_with_params(&model_path, WhisperContextParameters::default())
        .expect("model load");

    // Warm-up pass to prime caches.
    {
        let mut state = ctx.create_state().expect("state");
        let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
        params.set_language(Some("en"));
        params.set_n_threads(physical as i32);
        params.set_print_special(false);
        params.set_print_progress(false);
        params.set_print_realtime(false);
        state.full(params, &samples).expect("warmup");
    }

    let mut targets: Vec<i32> = vec![1, 2, 4, physical as i32, logical as i32];
    if logical >= 8 && !targets.contains(&8) {
        targets.push(8);
    }
    targets.sort();
    targets.dedup();

    // Snapshot the runtime Vulkan loader status once. The actual whisper
    // context above already initialised whichever backend it could; the
    // GPU panels below differ only in label and predicted-helper-pick.
    // The runtime RTF rows are produced by the same backend the warm-up
    // used.
    let vulkan_runtime_ok = cfg!(feature = "whisper-vulkan") && vulkan_loader_available();
    eprintln!(
        "[sweep] whisper-vulkan feature: {}, libvulkan resolvable at runtime: {}",
        cfg!(feature = "whisper-vulkan"),
        vulkan_runtime_ok
    );

    // Four panels: CPU and GPU axes for the predicted-helper-pick column,
    // crossed with AC and battery via MAGNOTIA_POWER_STATE_OVERRIDE.
    let panels = [
        ("AC, CPU", "ac", false),
        ("AC, GPU (Vulkan)", "ac", true),
        ("battery, CPU", "battery", false),
        ("battery, GPU (Vulkan)", "battery", true),
    ];

    for (label, power, gpu_offloaded_for_helper) in panels {
        env::set_var("MAGNOTIA_POWER_STATE_OVERRIDE", power);
        let helper_pick = inference_thread_count(Workload::Whisper, gpu_offloaded_for_helper);
        run_sweep_panel(
            label,
            helper_pick,
            &ctx,
            &samples,
            audio_secs,
            &targets,
        );
    }
    env::remove_var("MAGNOTIA_POWER_STATE_OVERRIDE");
}

fn run_sweep_panel(
    label: &str,
    helper_pick: usize,
    ctx: &WhisperContext,
    samples: &[f32],
    audio_secs: f64,
    targets: &[i32],
) {
    eprintln!("");
    eprintln!(
        "=== n_threads scaling: {label} (helper picks: {helper_pick}) ==="
    );
    eprintln!("n_threads | xc_time | RTF    | speedup_vs_1");
    eprintln!("----------|---------|--------|-------------");
    let mut baseline_dur: Option<f64> = None;
    for n in targets {
        // Two runs, take the min — best-case after L2/L3 warm.
        let mut best = f64::MAX;
        for _ in 0..2 {
            let mut state = ctx.create_state().expect("state");
            let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 1 });
            params.set_language(Some("en"));
            params.set_n_threads(*n);
            params.set_print_special(false);
            params.set_print_progress(false);
            params.set_print_realtime(false);
            let t = Instant::now();
            state.full(params, samples).expect("transcribe");
            let dur = t.elapsed().as_secs_f64();
            if dur < best {
                best = dur;
            }
        }
        let rtf = best / audio_secs;
        let speedup = baseline_dur.map(|b| b / best).unwrap_or(1.0);
        if baseline_dur.is_none() {
            baseline_dur = Some(best);
        }
        eprintln!(
            "{:>9} | {:>6.2}s | {:>6.3} | {:>6.2}x",
            n, best, rtf, speedup
        );
    }
}