cortiq-gateway/src/bench.rs

//! One-click local model benchmark: run `cortiq bench <model> --json` and keep
//! the parsed result per model, so the Models page can show real tok/s numbers
//! measured on this machine.

use serde::Serialize;
use std::collections::HashMap;
use std::process::Stdio;
use std::sync::{Arc, Mutex};

#[derive(Clone, Serialize)]
pub struct BenchRun {
    pub model: String, // file stem under models_dir
    pub state: String, // running | done | error
    pub log: Vec<String>,
    pub tok_s: Option<f64>,
    /// The full JSON object the benchmark printed, passed through untouched.
    pub raw: Option<serde_json::Value>,
    pub started: u64,
    pub finished: Option<u64>,
}

#[derive(Default)]
pub struct BenchStore {
    runs: Mutex<HashMap<String, BenchRun>>,
    busy: Mutex<bool>,
}

fn now() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// tok/s out of the benchmark's JSON: known field names first (steady decode
/// speed is the honest number), then any `*tok_s*` field as a fallback. The
/// old loose "contains tok + per" heuristic could grab `allocs_per_token`.
fn extract_tok_s(v: &serde_json::Value) -> Option<f64> {
    let obj = v.as_object()?;
    for want in [
        "decode_tok_s_steady",
        "decode_tok_s_incl_prefill",
        "decode_tok_s",
        "tok_s",
        "tok_per_s",
        "tokens_per_sec",
    ] {
        if let Some(n) = obj.get(want).and_then(|x| x.as_f64()) {
            return Some(n);
        }
    }
    obj.iter()
        .find(|(k, val)| {
            let k = k.to_ascii_lowercase();
            (k.contains("tok_s") || k.contains("tokens_per")) && val.as_f64().is_some()
        })
        .and_then(|(_, val)| val.as_f64())
}

/// Last parseable JSON object in mixed CLI output. Newer cortiq releases
/// pretty-print the `--json` result over many lines, older ones emit a single
/// line — this handles both, skipping any log noise around it.
fn last_json_object(text: &str) -> Option<serde_json::Value> {
    let mut end = text.len();
    while let Some(close) = text[..end].rfind('}') {
        let mut depth = 0i32;
        let mut start = None;
        for (i, c) in text[..=close].char_indices().rev() {
            match c {
                '}' => depth += 1,
                '{' => {
                    depth -= 1;
                    if depth == 0 {
                        start = Some(i);
                        break;
                    }
                }
                _ => {}
            }
        }
        if let Some(st) = start {
            if let Ok(v) = serde_json::from_str::<serde_json::Value>(&text[st..=close]) {
                if v.is_object() {
                    return Some(v);
                }
            }
        }
        end = close;
    }
    None
}

impl BenchStore {
    pub fn new() -> Arc<Self> {
        Arc::new(Self::default())
    }
    pub fn get(&self, model: &str) -> Option<BenchRun> {
        self.runs.lock().unwrap().get(model).cloned()
    }
    pub fn list(&self) -> Vec<BenchRun> {
        self.runs.lock().unwrap().values().cloned().collect()
    }
}

/// Kick off a benchmark for `<models_dir>/<model>.cmf`. One at a time — a
/// benchmark loads the whole model and owns the machine while it runs.
pub fn start(
    store: Arc<BenchStore>,
    cmf: &crate::config::CmfCfg,
    model: String,
) -> Result<(), String> {
    let path = std::path::Path::new(&cmf.models_dir).join(format!("{model}.cmf"));
    if !path.exists() {
        return Err(format!("model file not found: {}", path.display()));
    }
    {
        let mut busy = store.busy.lock().unwrap();
        if *busy {
            return Err("a benchmark is already running — wait for it to finish".into());
        }
        *busy = true;
    }
    store.runs.lock().unwrap().insert(
        model.clone(),
        BenchRun {
            model: model.clone(),
            state: "running".into(),
            log: Vec::new(),
            tok_s: None,
            raw: None,
            started: now(),
            finished: None,
        },
    );
    let bin = cmf.cortiq_bin.clone();
    let gpu = cmf.gpu;
    tokio::spawn(async move {
        let out = tokio::process::Command::new(&bin)
            .arg("bench")
            .arg(&path)
            .arg("--json")
            .env("CMF_GPU", if gpu { "1" } else { "0" })
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .output()
            .await;
        let mut g = store.runs.lock().unwrap();
        let run = g.get_mut(&model).unwrap();
        run.finished = Some(now());
        match out {
            Ok(o) => {
                let text = format!(
                    "{}\n{}",
                    String::from_utf8_lossy(&o.stdout),
                    String::from_utf8_lossy(&o.stderr)
                );
                // keep a readable tail; the JSON line is parsed separately
                run.log = text
                    .lines()
                    .map(|l| l.trim())
                    .filter(|l| !l.is_empty() && !l.contains("Metal GPU path"))
                    .map(String::from)
                    .collect::<Vec<_>>();
                let n = run.log.len();
                if n > 30 {
                    run.log.drain(0..n - 30);
                }
                let parsed = last_json_object(&text);
                match parsed {
                    Some(v) if o.status.success() => {
                        run.tok_s = extract_tok_s(&v);
                        run.raw = Some(v);
                        run.state = "done".into();
                    }
                    _ => run.state = "error".into(),
                }
            }
            Err(e) => {
                run.log.push(format!("✗ failed to start cortiq bench: {e}"));
                run.state = "error".into();
            }
        }
        drop(g);
        *store.busy.lock().unwrap() = false;
    });
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn tok_s_from_common_keys() {
        let v = serde_json::json!({"tok_s": 42.5, "allocs_per_token": 0});
        assert_eq!(extract_tok_s(&v), Some(42.5));
        let v = serde_json::json!({"decode_tokens_per_sec": 31.0});
        assert_eq!(extract_tok_s(&v), Some(31.0));
    }

    #[test]
    fn steady_decode_wins_and_allocs_never_matches() {
        let v = serde_json::json!({
            "allocs_per_token": 769.4,
            "decode_tok_s_incl_prefill": 9.98,
            "decode_tok_s_steady": 11.86,
            "prefill_tok_s": 25.5,
        });
        assert_eq!(extract_tok_s(&v), Some(11.86));
        let v = serde_json::json!({"allocs_per_token": 769.4, "seq_len": 140});
        assert_eq!(extract_tok_s(&v), None);
    }

    #[test]
    fn json_found_single_line_and_pretty() {
        let single = "noise\n{\"tok_s\": 5.0}\nmore noise";
        assert_eq!(last_json_object(single).unwrap()["tok_s"], 5.0);
        let pretty = "INFO loading\n{\n  \"decode_tok_s_steady\": 11.86,\n  \"seq_len\": 140\n}\nINFO done {not json";
        assert_eq!(
            last_json_object(pretty).unwrap()["decode_tok_s_steady"],
            11.86
        );
        assert!(last_json_object("no braces here").is_none());
    }
}