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::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::{Arc, Mutex};

use anyhow::Result;

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,
    /// Skip the context-depth suite (16K..256K) — run only the main pass
    /// (scenarios + the typical-depth perf rows).
    pub no_depth: 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>,
    /// When set, the tool-eval-bench invocation runs on this host over SSH
    /// (the local machine only orchestrates), keeping the bench client's own
    /// network (e.g. a WiFi laptop) off the measurement path. The tokenizer
    /// is uploaded to the host; the report + run package land locally. The
    /// host must have `tool-eval-bench` on PATH and reach the serve.
    pub ssh_host: 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()
}

/// 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 (2026-09-09). The wrapper is
/// macOS-only (BSD `script`); elsewhere 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) that once leaked
/// into the saved .log (glm53-flash-eugr-2x, 2026-09-12).
/// `TTY_COMPATIBLE=0` (rich's own knob) renders non-interactively (bars
/// silent, one final frame) while the pty keeps the deadlock workaround;
/// `NO_COLOR=1` covers non-rich output.
fn run_tool(
    bin: &str,
    args: &[String],
    cwd: &Path,
    log_path: &Path,
    append_log: bool,
    ssh_host: Option<&str>,
) -> 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)?
    }));
    // The command: local (the TTY wrapper or direct exec) or remote (SSH to
    // the host — a pty via `ssh -t` keeps the rich Live display from
    // deadlocking on the non-tty pipe; `stty -onlcr` keeps the log LF-only;
    // TTY_COMPATIBLE/NO_COLOR match the local run). For the remote host the
    // `cwd` is a remote path embedded in the shell string, so the local
    // current_dir/env are skipped.
    let mut cmd = if let Some(host) = ssh_host {
        crate::bench_remote::remote_command(host, cwd, bin, args)
    } else if cfg!(target_os = "macos")
        && 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
    };
    if ssh_host.is_none() {
        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 || crate::bench_io::pump(stdout_pipe, true, f1));
    let t2 = std::thread::spawn(move || crate::bench_io::pump(stderr_pipe, false, file));
    let status = child.wait()?;
    t1.join().ok();
    t2.join().ok();
    Ok(status.code().unwrap_or(-1))
}

/// 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 log_path = reports_dir.join(format!("{stamp}_{label}.log"));
    let ssh_host = a.ssh_host.as_deref();
    // The local scratch dir (the report reads the md + json from here). For
    // the remote host the tool writes to a remote scratch dir; its md/json are
    // downloaded here after the run so the local report generation is unchanged.
    let scratch = dirs::home_dir()
        .ok_or_else(|| anyhow::anyhow!("no home dir"))?
        .join(".cache")
        .join("sm121")
        .join("bench")
        .join(format!("{stamp}_{label}"));
    // The scratch dir the tool writes to (local by default, or the remote host
    // when `--ssh-host` is set — the bench client's network then stays off the
    // measurement path).
    let tool_scratch = if ssh_host.is_some() {
        PathBuf::from(format!("/tmp/sm121-bench-{stem}"))
    } else {
        scratch.clone()
    };
    // The local json path (the report reads it).
    let json_path = reports_dir.join(format!("{stamp}_{label}.json"));
    // The json path the tool writes to (local by default, or the remote host).
    let tool_json = if ssh_host.is_some() {
        tool_scratch.join(format!("{stem}.json"))
    } else {
        json_path.clone()
    };
    let bin = "tool-eval-bench".to_string();
    // The tokenizer path the tool uses (local by default, or the remote
    // scratch dir when `--ssh-host` is set — the tokenizer is uploaded there).
    let tool_tokenizer = if ssh_host.is_some() {
        a.tokenizer.as_ref().map(|tok| {
            let name = tok.rsplit('/').next().unwrap_or("tokenizer");
            tool_scratch.join(name).to_string_lossy().to_string()
        })
    } else {
        a.tokenizer.clone()
    };
    let depth_scratch = tool_scratch.with_file_name(format!("{stem}_depth"));
    let depth_args =
        crate::bench_depth::depth_args(&base_url, &model, a.seed, tool_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(),
        // Depth rows must decode a WARM prefix: llama-benchy's default
        // branch concatenates context+prompt into one 67K request whose
        // serialized dual prefills measured 13.7/5.6/2.1 t/s at c2 — an
        // admission artifact, not the engine (diagnosis 2026-09-29,
        // wire capture + this flag flipping c2 13.7 → 52.1 in one run).
        "--benchy-args=--enable-prefix-caching".into(),
        "--concurrency".into(),
        a.concurrency.clone(),
        "--json-file".into(),
        tool_json.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) = &tool_tokenizer {
        args.push("--tokenizer".into());
        args.push(tok.clone());
    }

    if dry_run {
        println!("bench: cwd {}", tool_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(());
    }

    if let Some(host) = ssh_host {
        crate::bench_remote::prepare(host, &tool_scratch, &a.tokenizer)?;
    } else {
        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, &tool_scratch, &log_path, false, ssh_host)?;
    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.
    // `--no-depth` skips the whole phase (quick-grasp campaigns).
    if !a.no_depth {
        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, ssh_host)?;
        if dcode != 0 {
            eprintln!(
                "bench: WARNING depth suite exited with {dcode} — main report unaffected ({})",
                log_path.display()
            );
        }
    } else {
        eprintln!("bench: depth suite skipped (--no-depth)");
    }

    // Remote host: download the json + the md from the remote scratch dir so
    // the local report generation is unchanged.
    if let Some(host) = ssh_host {
        crate::bench_remote::collect(host, &tool_json, &json_path, &tool_scratch, &scratch)?;
    }

    let md_src = crate::bench_io::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 || crate::bench_io::pump(out, false, f1));
        let t2 = std::thread::spawn(move || crate::bench_io::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,
            None,
        )
        .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, None).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 = "name = \"t\"\ndocker_image = \"x:1\"\nmodel = \"org/m\"\nargs = [\"--served-model-name\", \"glm53\"]\n";
        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());
    }
}