launcher/src/bench_target.rs

//! Bench target resolution + the tool runner, split out of `bench.rs` at the
//! 500-line cap: the served model name, the host node (`--node` / rank 0), a
//! filesystem-safe label, and the pty-wrapped `tool-eval-bench` invocation
//! that tees its output into the log file.

use crate::config::Recipe;
use anyhow::Result;
use std::fs;
use std::path::Path;
use std::process::Stdio;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;

pub(crate) fn served_model(recipe: &Recipe) -> Option<String> {
    let args = recipe.args.to_vec();
    // `--served-model-name` is the vLLM/SGLang spelling; TensorFold's HTTP app
    // takes `--name` instead (its 2x 1017 lane serves `deepseek-v4.1-flash`).
    ["--served-model-name", "--name"].iter().find_map(|flag| {
        args.iter()
            .position(|s| s == flag || s.starts_with(&format!("{flag}=")))
            .and_then(|i| match args.get(i) {
                Some(a) if a.contains('=') => a.split_once('=').map(|(_, v)| v.to_string()),
                _ => args.get(i + 1).cloned(),
            })
    })
}

/// Bench target node: the explicit `--node` override, else rank 0 (falling
/// back to the first node). `--node` matters for serves pinned with
/// `serve --node <n>` — the run's head may not be the cluster's rank 0
/// (e.g. a solo Bolt on dgx3 inside `4nodes-switch` while dgx1 serves
/// something else entirely).
pub(crate) fn pick_head<'a>(
    nodes: &'a [crate::config::Node],
    cluster_name: &str,
    node: Option<&str>,
) -> anyhow::Result<&'a crate::config::Node> {
    match node {
        Some(name) => {
            let bare = name.strip_suffix(".lan").unwrap_or(name);
            nodes
                .iter()
                .find(|n| n.name.strip_suffix(".lan").unwrap_or(&n.name) == bare)
                .ok_or_else(|| {
                    anyhow::anyhow!("node '{}' is not part of cluster '{}'", name, cluster_name)
                })
        }
        None => nodes
            .iter()
            .find(|n| n.rank == 0)
            .or_else(|| nodes.first())
            .ok_or_else(|| anyhow::anyhow!("cluster '{}' has no nodes", cluster_name)),
    }
}

/// Keep the label filesystem-safe.
pub(crate) 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.
pub(crate) fn run_tool(
    bin: &str,
    args: &[String],
    cwd: &Path,
    log_path: &Path,
    append_log: bool,
    ssh_host: Option<&str>,
    cancel: Option<&Arc<AtomicBool>>,
) -> 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));
    // Bench cancellation: poll the job's flag; on request kill the harness
    // child (the SSH child for a remote run, so the remote command drops too).
    let status = loop {
        if cancel.is_some_and(|c| c.load(Ordering::Relaxed)) {
            let _ = child.kill();
            break child.wait()?;
        }
        if let Some(s) = child.try_wait()? {
            break s;
        }
        thread::sleep(Duration::from_millis(200));
    };
    t1.join().ok();
    t2.join().ok();
    Ok(status.code().unwrap_or(-1))
}