| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157 |
- //! 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::{Arc, Mutex};
- 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>,
- ) -> 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))
- }
|