launcher/src/bench.rs

//! `sm121 bench` — run `tool-eval-bench` against a live serve and land the
//! reports in `recipes/<lane>/reports/`. The tool needs a scratch cwd (it
//! writes `runs/<date>/<run>.md` relative to cwd), an absolute `--json-file`
//! path, and the `.md` copied into the lane — this command does the
//! babysitting. Campaign convention (seed 42, hardmode, perf, c1,2) is the
//! default; `sm121 bench -d` prints the exact command.

use std::fs;
use std::path::PathBuf;

use anyhow::Result;

use crate::bench_target::{pick_head, run_tool, sanitize, served_model};
use crate::config::{Config, load_recipe, select_cluster_by_name};

/// Fields of the `sm121 bench` subcommand (mirrored by the clap variant in
/// `cli.rs`). Served name/port come from the recipe; the tool is on PATH.
#[derive(Clone)]
pub struct Args {
    pub recipe: String,
    pub cluster: String,
    pub label: Option<String>,
    pub concurrency: String,
    pub seed: u64,
    /// Scenario streams run concurrently (tool `--parallel`; None = serial).
    pub parallel: Option<u32>,
    /// Skip the scenario phase (tool `--perf-only`).
    pub perf_only: bool,
    /// Run the context-depth suite (16K..256K) after the main pass — off by
    /// default (a full sweep takes ~45 min on a 2x lane).
    pub 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>,
    /// Bench target node within the cluster (default: rank 0). Use when the
    /// serve was launched with `serve --node <n>` — a pinned/solo run may
    /// live on a non-rank-0 node, and benching the wrong head measures a
    /// different (or no) engine.
    pub node: Option<String>,
    /// Disable thinking/reasoning in the harness (tool `--no-think`, sets
    /// `enable_thinking=false`). Needed for thinking-by-default models (Ling,
    /// GLM-4.5-style) whose long default reasoning blows the per-turn timeout
    /// and excludes most scenarios.
    pub no_think: bool,
    /// Per-turn request timeout seconds forwarded to the harness (tool
    /// `--timeout`). Raise above the 120 s default for slow/thinking models.
    pub timeout: Option<u64>,
    /// Job cancellation flag (web bench window). `run_tool` polls it and
    /// kills the harness child; None for the CLI.
    pub cancel: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
}

/// Served model name from the recipe's `--served-model-name` arg.
/// 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 = pick_head(&cluster.nodes, &cluster.name, a.node.as_deref())?;
    // FLARE serves on its own fixed API port (`[flare].http_port`), not the
    // cluster default the docker lanes use.
    let port = if recipe.engine == "flare" {
        crate::native::flare::http_port(&recipe)
    } else {
        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();
    // Label: required for quality runs, optional for perf-only (default `perf`).
    let label = match (&a.label, a.perf_only) {
        (Some(l), _) => sanitize(l),
        (None, true) => String::from("perf"),
        (None, false) => {
            anyhow::bail!("--label is required for a quality run (the campaign record)")
        }
    };
    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/sm121/bench")
        .join(&stem);
    // 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(),
        // Warm-prefix decode: without it llama-benchy's default branch
        // serializes dual 67K prefills (c2 collapses to 13.7 t/s — an
        // admission artifact, not the engine; 2026-09-29 diagnosis).
        "--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 a.no_think {
        args.push("--no-think".into());
    }
    if let Some(t) = a.timeout {
        args.push("--timeout".into());
        args.push(t.to_string());
    }

    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 -> {} (run package: {}.recipe.toml + {}.manifest.json)",
            reports_dir.display(),
            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,
        a.cancel.as_ref(),
    )?;
    if code != 0 {
        if a.cancel
            .as_ref()
            .is_some_and(|c| c.load(std::sync::atomic::Ordering::Relaxed))
        {
            anyhow::bail!("bench: cancelled");
        }
        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.
    // Opt-in (`--depth`): a full sweep takes ~45 min on a 2x lane.
    if a.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,
            a.cancel.as_ref(),
        )?;
        if dcode != 0 {
            if a.cancel
                .as_ref()
                .is_some_and(|c| c.load(std::sync::atomic::Ordering::Relaxed))
            {
                anyhow::bail!("bench: cancelled in depth suite");
            }
            eprintln!(
                "bench: WARNING depth suite exited with {dcode} — main report unaffected ({})",
                log_path.display()
            );
        }
    } else {
        eprintln!("bench: depth suite skipped (opt-in: --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,
            !a.perf_only,
        )?;
    }

    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 crate::config::{Node, Recipe};
    use std::process::{Command, Stdio};
    use std::sync::{Arc, Mutex};

    fn node(name: &str, rank: u32) -> Node {
        Node {
            name: name.to_string(),
            ssh_ip: "192.0.2.1".into(),
            ib_ip: None,
            rank,
            host_key: None,
            entity: None,
        }
    }

    #[test]
    fn pick_head_defaults_to_rank0() {
        let nodes = vec![node("a.lan", 2), node("b.lan", 0), node("c.lan", 1)];
        assert_eq!(pick_head(&nodes, "c", None).unwrap().name, "b.lan");
    }

    #[test]
    fn pick_head_node_override_tolerates_lan_suffix() {
        let nodes = vec![node("a.lan", 2), node("b.lan", 0)];
        assert_eq!(pick_head(&nodes, "c", Some("a")).unwrap().name, "a.lan");
        assert_eq!(pick_head(&nodes, "c", Some("b.lan")).unwrap().name, "b.lan");
        assert!(pick_head(&nodes, "c", Some("zzz")).is_err());
    }

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