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::SystemTime;

use anyhow::Result;

use crate::ansi::strip_ansi_step;
use crate::config::{Config, Recipe, load_recipe, 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,
    /// Scenario streams run concurrently (tool `--parallel`). None = the
    /// tool default (serial), preserving earlier lanes' comparability.
    pub parallel: Option<u32>,
    /// Skip the scenario phase (tool `--perf-only`).
    pub perf_only: bool,
    /// Local tokenizer source for llama-benchy (tool `--tokenizer`): a
    /// tokenizer.json or a dir holding one. Needed when the served model
    /// id is not resolvable in the local HF cache (custom export lanes).
    pub tokenizer: Option<String>,
}

/// 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()
}

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

/// Tee one output pipe to the console and the log file.
///
/// The console copy keeps the child's ANSI colors; the file copy is
/// stripped (state machine below) so the saved .log stays readable in
/// `sm121 report` and under grep/less.
///
/// 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];
    let mut carry = Vec::new(); // trailing partial escape across chunks
    let mut clean = Vec::new();
    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();
        }
        carry.extend_from_slice(&buf[..n]);
        strip_ansi_step(&mut carry, &mut clean);
        if !clean.is_empty() {
            if let Ok(mut f) = file.lock() {
                f.write_all(&clean).ok();
            }
            clean.clear();
        }
    }
}

/// Run the bench tool in `cwd`, teeing stdout+stderr to the console and the
/// log file; returns the exit code.
///
/// `append_log = false` truncates the log (main campaign phase); `true`
/// appends (the depth-suite phase shares one log so its rows land in the
/// report's Throughput table).
///
/// 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.
///
/// The pty makes the tool think it is interactive — rich redraws its
/// spinner/progress bars in place (\r soup on one line) and that soup
/// leaked into the saved .log, even shadowing real parser cells
/// (glm53-flash-eugr-2x, 2026-09-12). `TTY_COMPATIBLE=0` is rich's own
/// knob for this: the console renders non-interactively (progress bars
/// stay silent during the run, one final frame at the end) while the pty
/// keeps the deadlock workaround intact. `NO_COLOR=1` covers non-rich
/// output. Verified locally under `script`: 1.3 KB/8 CR of soup → 141
/// bytes/1 static final frame.
fn run_tool(
    bin: &str,
    args: &[String],
    cwd: &Path,
    log_path: &Path,
    append_log: bool,
) -> Result<i32> {
    let file = Arc::new(Mutex::new(if append_log {
        fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(log_path)?
    } else {
        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()
    {
        // Wrap in `sh` + `stty -onlcr`: the macOS pty's ONLCR would turn every
        // `\n` the tool writes into `\r\n` on disk, and the report viewer
        // then renders each line twice (blank line after every line).
        // `stty -onlcr` disables the post-processing at the source; `exec`
        // replaces the shell so the tool keeps the pty and the exit code.
        let mut c = std::process::Command::new("/usr/bin/script");
        c.arg("-q");
        c.arg("/dev/null");
        c.arg("sh");
        c.arg("-c");
        c.arg("stty -onlcr 2>/dev/null; exec \"$@\"");
        c.arg("sh"); // $0 inside the -c script
        c.arg(bin);
        c.args(args);
        c
    } else {
        let mut c = std::process::Command::new(bin);
        c.args(args);
        c
    };
    cmd.current_dir(cwd);
    cmd.env("TTY_COMPATIBLE", "0");
    cmd.env("NO_COLOR", "1");
    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 = crate::config::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 = crate::stamp::utc_stamp();
    let label = sanitize(&a.label);
    let stem = format!("{stamp}_{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 depth_scratch = scratch.with_file_name(format!("{stem}_depth"));
    let depth_args =
        crate::bench_depth::depth_args(&base_url, &model, a.seed, a.tokenizer.as_ref());

    let mut 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 let Some(p) = a.parallel {
        args.push("--parallel".into());
        args.push(p.to_string());
    }
    if a.perf_only {
        args.push("--perf-only".into());
    }
    if let Some(tok) = &a.tokenizer {
        args.push("--tokenizer".into());
        args.push(tok.clone());
    }

    if dry_run {
        println!("bench: cwd {}", scratch.display());
        println!("bench: {} {}", bin, args.join(" "));
        println!("bench: depth suite: {} {}", bin, depth_args.join(" "));
        println!("bench: reports -> {}", reports_dir.display());
        println!(
            "bench: run package -> {}.recipe.toml + {}.manifest.json",
            stem, stem
        );
        return Ok(());
    }

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

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

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

    // Phase 2 — context-depth suite (16K..256K @ c2, 2 repeats): appended
    // to the same log so `parse_perf_rows` picks its rows up too. A depth
    // failure must not cost the finished quality run — warn, keep going.
    eprintln!(
        "bench: depth suite: {} @ c{}, {} repeats{}",
        crate::bench_depth::DEPTHS,
        crate::bench_depth::CONCURRENCY,
        crate::bench_depth::RUNS,
        if a.tokenizer.is_some() {
            " (+tokenizer)"
        } else {
            ""
        }
    );
    fs::create_dir_all(&depth_scratch)?;
    let dcode = run_tool(&bin, &depth_args, &depth_scratch, &log_path, true)?;
    if dcode != 0 {
        eprintln!(
            "bench: WARNING depth suite exited with {dcode} — main report unaffected ({})",
            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)?;

    // The `--json-file` dispatch path drops the perf section from the md —
    // re-attach it from the llama-benchy table in the log.
    if crate::bench_report::ensure_throughput_section(&md_dest, &log_path).is_ok_and(|m| m) {
        eprintln!("bench: appended the perf table to the report md");
    }

    // Summary: final score from the JSON + the perf table from the log.
    // Perf-only runs have no quality score and the tool writes no JSON —
    // the run package is skipped for them (the campaign run is the record).
    let j = match a.perf_only {
        true => {
            eprintln!("bench: perf-only run (no quality score, no run package)");
            None
        }
        false => {
            let j: serde_json::Value = serde_json::from_str(&fs::read_to_string(&json_path)?)?;
            let sc = &j["scores"];
            let final_score = sc["final_score"].as_u64();
            if let Some(score) = final_score {
                let completion = sc["completion_rate"]
                    .as_f64()
                    .map(|v| v.to_string())
                    .unwrap_or_default();
                eprintln!(
                    "bench: final score {score} / {} ({} points), completion {completion}",
                    sc["max_points"], sc["total_points"]
                );
            } else {
                eprintln!("bench: perf-only run (no quality score)");
            }
            Some(j)
        }
    };
    for line in fs::read_to_string(&log_path)?.lines() {
        if line.contains("pp2048") && line.contains("tg128") {
            println!("{line}");
        }
    }
    if let Some(j) = &j {
        crate::bench_artifacts::write_run_package(&crate::bench_artifacts::RunPackage {
            config,
            cluster: &cluster,
            recipe: &recipe,
            recipe_file: &file,
            stem: stem.as_str(),
            reports_dir: &reports_dir,
            bench_args: &args,
            bench_json: j,
        })
        .await?;
    }
    eprintln!("bench: reports landed in {}", reports_dir.display());
    Ok(())
}

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

    /// 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");
    }

    /// The non-interactive env must reach the child through the `script`
    /// pty wrapper — rich reads TTY_COMPATIBLE from the child's env. Also:
    /// `stty -onlcr` must hold, so the saved log is LF-only (CRLF made the
    /// report viewer render every line twice).
    #[test]
    fn run_tool_passes_non_interactive_env() {
        let tmp = std::env::temp_dir().join(format!("sm121-bench-env-test-{}", std::process::id()));
        fs::create_dir_all(&tmp).unwrap();
        let log = tmp.join("env.log");
        let code = run_tool(
            "sh",
            &[
                "-c".to_string(),
                "echo TTY=$TTY_COMPATIBLE NO=$NO_COLOR".to_string(),
            ],
            &tmp,
            &log,
            false,
        )
        .unwrap();
        let body = fs::read_to_string(&log).unwrap();
        let _ = fs::remove_dir_all(&tmp);
        assert_eq!(code, 0, "exit code, log: {body:?}");
        assert!(body.contains("TTY=0"), "TTY_COMPATIBLE lost: {body:?}");
        assert!(body.contains("NO=1"), "NO_COLOR lost: {body:?}");
        assert!(!body.contains('\r'), "pty ONLCR leaked CRLF: {body:?}");
    }

    /// The append mode must not truncate the main phase's log — the depth
    /// suite shares it so one report carries both tables.
    #[test]
    fn run_tool_append_keeps_prior_output() {
        let tmp =
            std::env::temp_dir().join(format!("sm121-bench-append-test-{}", std::process::id()));
        fs::create_dir_all(&tmp).unwrap();
        let log = tmp.join("append.log");
        let run =
            |a: &str, append| run_tool("sh", &["-c".into(), a.into()], &tmp, &log, append).unwrap();
        assert_eq!(run("echo main", false), 0);
        assert!(fs::read_to_string(&log).unwrap().contains("main\n"));
        assert_eq!(run("echo depth", true), 0);
        let body = fs::read_to_string(&log).unwrap();
        let (i1, i2) = (body.find("main\n"), body.find("depth\n"));
        assert!(
            i1.is_some() && i2.is_some_and(|d| i1.unwrap() < d),
            "append truncated the log: {body:?}"
        );
        let _ = fs::remove_dir_all(&tmp);
    }

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