launcher/src/bench.rs

//! `sm121 bench` — run `tool-eval-bench` against a live serve and land the
//! reports in the repo.
//!
//! The bench tool needs babysitting this command does instead of the
//! operator: it must run in a scratch cwd (it writes
//! `runs/<date>/<run>.md` relative to cwd), needs an absolute `--json-file`
//! path, and the `.md` report must be copied into
//! `recipes/<lane>/reports/<UTC-date>_<label>.{md,json,log}` by hand.
//! Campaign convention (seed 42, hardmode, perf, concurrency 1,2) is the
//! default; `sm121 bench -d` prints the exact command.

use std::fs;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use anyhow::Result;

use crate::config::{
    Config, Recipe, default_config_dir, load_recipe, local_path, select_cluster_by_name,
};

/// Fields of the `sm121 bench` subcommand (mirrored by the clap variant in
/// `cli.rs`). Served model name and port come from the recipe
/// (`--served-model-name` in args / `port`); the bench binary is looked up on PATH.
#[derive(Clone)]
pub struct Args {
    pub recipe: String,
    pub cluster: String,
    pub label: String,
    pub concurrency: String,
    pub seed: u64,
}

/// How long to poll the API for readiness before giving up.
const READY_TIMEOUT_SECS: u64 = 900;
/// Readiness poll interval.
const READY_POLL_SECS: u64 = 5;

/// Locate the recipe file (same search order as [`load_recipe`]) so reports
/// land next to it.
fn recipe_file(recipe_name: &str, sources: &[String]) -> Option<PathBuf> {
    let mut paths: Vec<PathBuf> = vec![
        PathBuf::from(format!("recipes/{}/recipe.toml", recipe_name)),
        PathBuf::from(format!("{}/recipe.toml", recipe_name)),
    ];
    for src in sources {
        paths.push(local_path(src).join(recipe_name).join("recipe.toml"));
    }
    if let Some(dir) = default_config_dir() {
        paths.push(dir.join("recipes").join(recipe_name).join("recipe.toml"));
    }
    paths.into_iter().find(|p| p.is_file())
}

/// Served model name from the recipe's `--served-model-name` arg.
fn served_model(recipe: &Recipe) -> Option<String> {
    let args = recipe.args.to_vec();
    args.iter()
        .position(|s| s == "--served-model-name")
        .and_then(|i| args.get(i + 1))
        .cloned()
}

/// `YYYY-MM-DD_HHMMSS` in UTC (no chrono dependency).
fn utc_stamp() -> String {
    let secs: i64 = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0);
    let (h, m, s) = ((secs % 86400) / 3600, (secs % 3600) / 60, secs % 60);
    let (y, mo, d) = civil_date(secs / 86400);
    format!(
        "{y:04}-{mo:02}-{d:02}_{h:02}{m:02}{s:02}",
        y = y,
        mo = mo,
        d = d
    )
}

/// Days since 1970-01-01 → (year, month, day); Howard Hinnant's civil_from_days.
fn civil_date(z: i64) -> (i64, i64, i64) {
    let z = z + 719_468;
    let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
    let doe = z - era * 146_097;
    let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let mo = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if mo <= 2 { y + 1 } else { y };
    (y, mo, d)
}

/// Keep the label filesystem-safe.
fn sanitize(label: &str) -> String {
    label
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
                c
            } else {
                '-'
            }
        })
        .collect()
}

/// Poll `GET <base>/models` until it answers (engine up).
async fn wait_ready(base_url: &str) -> Result<()> {
    let client = reqwest::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()?;
    let deadline = Instant::now() + Duration::from_secs(READY_TIMEOUT_SECS);
    let mut since_last_note = Instant::now();
    loop {
        match client.get(format!("{base_url}/models")).send().await {
            Ok(r) if r.status().is_success() => return Ok(()),
            Ok(_) => {}
            Err(_) => {}
        }
        if since_last_note.elapsed() >= Duration::from_secs(30) {
            since_last_note = Instant::now();
            eprintln!(
                "bench: waiting for API at {base_url} (up to {} s)...",
                READY_TIMEOUT_SECS
            );
        }
        if Instant::now() > deadline {
            anyhow::bail!("bench: API at {base_url} not ready after {READY_TIMEOUT_SECS} s");
        }
        tokio::time::sleep(Duration::from_secs(READY_POLL_SECS)).await;
    }
}

/// Tee one output pipe to the console and the log file.
/// Tee one output pipe to the console and the log file.
///
/// The read buffer must not be empty: `read()` on a zero-length slice
/// returns 0 (interpreted as EOF), so an empty `Vec` would silently sever
/// the pipe — the child then dies on its first stderr write with EPIPE
/// (silent exit 120 in CPython).
fn pump<R: Read>(mut reader: R, to_console: bool, file: Arc<Mutex<fs::File>>) {
    let mut buf = vec![0u8; 8192];
    loop {
        let n = match reader.read(&mut buf) {
            Ok(n) => n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
            Err(e) => panic!("reading bench output pipe: {e}"),
        };
        if n == 0 {
            break;
        }
        if to_console {
            std::io::stdout().write_all(&buf[..n]).ok();
        }
        if let Ok(mut f) = file.lock() {
            f.write_all(&buf[..n]).ok();
        }
    }
}

/// Run the bench tool in `cwd`, teeing stdout+stderr to the console and the
/// log file; returns the exit code.
///
/// The child runs under a pty (`script -q /dev/null`) instead of raw pipes:
/// tool-eval-bench's rich Live display deadlocks on a non-tty stdout pipe
/// right after the perf phase (deterministic on the 2x lane; 4x runs got
/// lucky) while it runs clean on a tty, exactly as in an interactive
/// terminal (2026-09-09). `script` is BSD/macOS form (the launcher is a
/// Mac tool); if it is absent we fall back to the direct exec.
fn run_tool(bin: &str, args: &[String], cwd: &Path, log_path: &Path) -> Result<i32> {
    let file = Arc::new(Mutex::new(fs::File::create(log_path)?));
    let mut cmd = if std::process::Command::new("/usr/bin/script")
        .arg("-q")
        .arg("/dev/null")
        .arg("true")
        .output()
        .is_ok()
    {
        let mut c = std::process::Command::new("/usr/bin/script");
        c.args(["-q", "/dev/null", bin]);
        c.args(args);
        c
    } else {
        let mut c = std::process::Command::new(bin);
        c.args(args);
        c
    };
    cmd.current_dir(cwd);
    cmd.stdout(Stdio::piped());
    cmd.stderr(Stdio::piped());
    let mut child = cmd.spawn()?;
    let stdout_pipe = child.stdout.take().expect("stdout pipe");
    let stderr_pipe = child.stderr.take().expect("stderr pipe");
    let f1 = file.clone();
    let t1 = std::thread::spawn(move || pump(stdout_pipe, true, f1));
    let t2 = std::thread::spawn(move || pump(stderr_pipe, false, file));
    let status = child.wait()?;
    t1.join().ok();
    t2.join().ok();
    Ok(status.code().unwrap_or(-1))
}

/// Newest `.md` under `scratch/runs/` (the tool's `runs/<date>/<run>.md`).
fn latest_md(scratch: &Path) -> Option<PathBuf> {
    let runs = scratch.join("runs");
    let mut best: Option<(std::time::SystemTime, PathBuf)> = None;
    let mut stack = vec![runs];
    while let Some(dir) = stack.pop() {
        let entries = fs::read_dir(dir).ok()?;
        for e in entries.flatten() {
            let p = e.path();
            if p.is_dir() {
                stack.push(p);
            } else if p.extension().is_some_and(|x| x == "md") {
                let t = e
                    .metadata()
                    .and_then(|m| m.modified())
                    .ok()
                    .unwrap_or(SystemTime::UNIX_EPOCH);
                if best.as_ref().is_none_or(|(bt, _)| t > *bt) {
                    best = Some((t, p));
                }
            }
        }
    }
    best.map(|(_, p)| p)
}

/// Run the bench and land `recipes/<lane>/reports/<stamp>_<label>.{md,json,log}`.
pub async fn run(config: &Config, a: &Args, dry_run: bool) -> Result<()> {
    let recipe = load_recipe(&a.recipe, &config.recipe_sources).await?;
    let file = recipe_file(&a.recipe, &config.recipe_sources)
        .ok_or_else(|| anyhow::anyhow!("recipe file not found for '{}'", a.recipe))?
        .canonicalize()?;
    let cluster = select_cluster_by_name(config, &a.cluster)?;
    let head = match cluster.nodes.iter().find(|n| n.rank == 0) {
        Some(n) => n,
        None => cluster
            .nodes
            .first()
            .ok_or_else(|| anyhow::anyhow!("cluster '{}' has no nodes", cluster.name))?,
    };
    let port = recipe.port;
    let model = served_model(&recipe).ok_or_else(|| {
        anyhow::anyhow!(
            "recipe '{}' has no --served-model-name in args — add it to the recipe",
            a.recipe
        )
    })?;
    let base_url = format!("http://{}:{}/v1", head.ssh_ip, port);
    let stamp = utc_stamp();
    let label = sanitize(&a.label);
    let reports_dir = file
        .parent()
        .ok_or_else(|| anyhow::anyhow!("no parent dir for {}", file.display()))?
        .join("reports");
    let json_path = reports_dir.join(format!("{stamp}_{label}.json"));
    let log_path = reports_dir.join(format!("{stamp}_{label}.log"));
    let scratch = dirs::home_dir()
        .ok_or_else(|| anyhow::anyhow!("no home dir"))?
        .join(".cache")
        .join("sm121")
        .join("bench")
        .join(format!("{stamp}_{label}"));
    let bin = "tool-eval-bench".to_string();

    let args = vec![
        "--base-url".into(),
        base_url.clone(),
        "--model".into(),
        model.clone(),
        "--seed".into(),
        a.seed.to_string(),
        "--hardmode".into(),
        "--perf".into(),
        "--concurrency".into(),
        a.concurrency.clone(),
        "--json-file".into(),
        json_path.to_string_lossy().to_string(),
    ];

    if dry_run {
        println!("bench: cwd {}", scratch.display());
        println!("bench: {} {}", bin, args.join(" "));
        println!("bench: reports -> {}", reports_dir.display());
        return Ok(());
    }

    fs::create_dir_all(&scratch)?;
    fs::create_dir_all(&reports_dir)?;

    eprintln!("bench: waiting for {base_url} ...");
    wait_ready(&base_url).await?;
    eprintln!("bench: API ready — running {} (seed {})", bin, a.seed);

    let code = run_tool(&bin, &args, &scratch, &log_path)?;
    if code != 0 {
        anyhow::bail!(
            "bench: tool exited with {code} (log: {})",
            log_path.display()
        );
    }

    let md_src = latest_md(&scratch)
        .ok_or_else(|| anyhow::anyhow!("no report .md under {}", scratch.join("runs").display()))?;
    let md_dest = reports_dir.join(format!("{stamp}_{label}.md"));
    fs::copy(&md_src, &md_dest)?;

    // Summary: final score from the JSON + the perf table from the log.
    let j: serde_json::Value = serde_json::from_str(&fs::read_to_string(&json_path)?)?;
    let sc = &j["scores"];
    let completion = sc["completion_rate"]
        .as_f64()
        .map(|v| v.to_string())
        .unwrap_or_default();
    eprintln!(
        "bench: final score {} / {} ({} points), completion {completion}",
        sc["final_score"], sc["max_points"], sc["total_points"]
    );
    for line in fs::read_to_string(&log_path)?.lines() {
        if line.contains("pp2048") && line.contains("tg128") {
            println!("{line}");
        }
    }
    eprintln!("bench: reports landed in {}", reports_dir.display());
    Ok(())
}

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

    #[test]
    fn stamp_format() {
        let s = utc_stamp();
        assert_eq!(s.len(), 17, "{s}");
        assert_eq!(&s[4..5], "-");
        assert_eq!(&s[7..8], "-");
        assert_eq!(&s[10..11], "_");
    }

    #[test]
    fn civil_date_known_values() {
        assert_eq!(civil_date(0), (1970, 1, 1));
        assert_eq!(civil_date(19_723), (2024, 1, 1));
        assert_eq!(civil_date(20_818), (2026, 12, 31));
    }

    /// Regression: the pump must not treat a not-yet-written pipe as EOF.
    /// `read()` into a zero-length buffer returns 0, which the old pump
    /// treated as EOF — closing the pipe and killing the child on its first
    /// stderr write with EPIPE (silent exit 120 in CPython).
    #[test]
    fn pump_keeps_pipe_open_until_child_exits() {
        let tmp =
            std::env::temp_dir().join(format!("sm121-bench-pump-test-{}", std::process::id()));
        fs::create_dir_all(&tmp).unwrap();
        let log = tmp.join("pump.log");

        let mut child = Command::new("sh")
            .args(["-c", "sleep 0.3; echo late-output >&2"])
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .unwrap();
        let out = child.stdout.take().unwrap();
        let err = child.stderr.take().unwrap();
        let f = Arc::new(Mutex::new(fs::File::create(&log).unwrap()));
        let f1 = f.clone();
        let t1 = std::thread::spawn(move || pump(out, false, f1));
        let t2 = std::thread::spawn(move || pump(err, false, f));
        let status = child.wait().unwrap();
        t1.join().unwrap();
        t2.join().unwrap();
        let body = fs::read_to_string(&log).unwrap();
        assert_eq!(status.code(), Some(0), "child must not hit EPIPE: {body}");
        assert!(body.contains("late-output"), "piped stderr lost: {body:?}");
        let _ = fs::remove_dir_all(&tmp);
    }

    #[test]
    fn sanitize_replaces_bad_chars() {
        assert_eq!(sanitize("arm B: 2x"), "arm-B--2x");
        assert_eq!(sanitize("ok_name-2.v3"), "ok_name-2.v3");
    }

    #[test]
    fn served_model_parsed_from_args() {
        let content = fs::read_to_string("recipes/glm53-flash-r27-spark/recipe.toml")
            .expect("repo recipe for test");
        let r: Recipe = toml::from_str(&content).expect("recipe toml");
        assert_eq!(served_model(&r).as_deref(), Some("glm53"));
        let mut r2 = r.clone();
        r2.args = crate::pre_serve::RecipeArgs::List(vec!["--trust-remote-code".into()]);
        assert!(served_model(&r2).is_none());
    }
}