bench_target.rs 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. //! Bench target resolution + the tool runner, split out of `bench.rs` at the
  2. //! 500-line cap: the served model name, the host node (`--node` / rank 0), a
  3. //! filesystem-safe label, and the pty-wrapped `tool-eval-bench` invocation
  4. //! that tees its output into the log file.
  5. use crate::config::Recipe;
  6. use anyhow::Result;
  7. use std::fs;
  8. use std::path::Path;
  9. use std::process::Stdio;
  10. use std::sync::{Arc, Mutex};
  11. pub(crate) fn served_model(recipe: &Recipe) -> Option<String> {
  12. let args = recipe.args.to_vec();
  13. // `--served-model-name` is the vLLM/SGLang spelling; TensorFold's HTTP app
  14. // takes `--name` instead (its 2x 1017 lane serves `deepseek-v4.1-flash`).
  15. ["--served-model-name", "--name"].iter().find_map(|flag| {
  16. args.iter()
  17. .position(|s| s == flag || s.starts_with(&format!("{flag}=")))
  18. .and_then(|i| match args.get(i) {
  19. Some(a) if a.contains('=') => a.split_once('=').map(|(_, v)| v.to_string()),
  20. _ => args.get(i + 1).cloned(),
  21. })
  22. })
  23. }
  24. /// Bench target node: the explicit `--node` override, else rank 0 (falling
  25. /// back to the first node). `--node` matters for serves pinned with
  26. /// `serve --node <n>` — the run's head may not be the cluster's rank 0
  27. /// (e.g. a solo Bolt on dgx3 inside `4nodes-switch` while dgx1 serves
  28. /// something else entirely).
  29. pub(crate) fn pick_head<'a>(
  30. nodes: &'a [crate::config::Node],
  31. cluster_name: &str,
  32. node: Option<&str>,
  33. ) -> anyhow::Result<&'a crate::config::Node> {
  34. match node {
  35. Some(name) => {
  36. let bare = name.strip_suffix(".lan").unwrap_or(name);
  37. nodes
  38. .iter()
  39. .find(|n| n.name.strip_suffix(".lan").unwrap_or(&n.name) == bare)
  40. .ok_or_else(|| {
  41. anyhow::anyhow!("node '{}' is not part of cluster '{}'", name, cluster_name)
  42. })
  43. }
  44. None => nodes
  45. .iter()
  46. .find(|n| n.rank == 0)
  47. .or_else(|| nodes.first())
  48. .ok_or_else(|| anyhow::anyhow!("cluster '{}' has no nodes", cluster_name)),
  49. }
  50. }
  51. /// Keep the label filesystem-safe.
  52. pub(crate) fn sanitize(label: &str) -> String {
  53. label
  54. .chars()
  55. .map(|c| {
  56. if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
  57. c
  58. } else {
  59. '-'
  60. }
  61. })
  62. .collect()
  63. }
  64. /// Run the bench tool in `cwd`, teeing stdout+stderr to the console and the
  65. /// log file; returns the exit code.
  66. ///
  67. /// `append_log = false` truncates the log (main campaign phase); `true`
  68. /// appends (the depth-suite phase shares one log so its rows land in the
  69. /// report's Throughput table).
  70. ///
  71. /// The child runs under a pty (`script -q /dev/null`) instead of raw pipes:
  72. /// tool-eval-bench's rich Live display deadlocks on a non-tty stdout pipe
  73. /// right after the perf phase (deterministic on the 2x lane; 4x runs got
  74. /// lucky) while it runs clean on a tty (2026-09-09). The wrapper is
  75. /// macOS-only (BSD `script`); elsewhere we fall back to the direct exec.
  76. ///
  77. /// The pty makes the tool think it is interactive — rich redraws its
  78. /// spinner/progress bars in place (\r soup on one line) that once leaked
  79. /// into the saved .log (glm53-flash-eugr-2x, 2026-09-12).
  80. /// `TTY_COMPATIBLE=0` (rich's own knob) renders non-interactively (bars
  81. /// silent, one final frame) while the pty keeps the deadlock workaround;
  82. /// `NO_COLOR=1` covers non-rich output.
  83. pub(crate) fn run_tool(
  84. bin: &str,
  85. args: &[String],
  86. cwd: &Path,
  87. log_path: &Path,
  88. append_log: bool,
  89. ssh_host: Option<&str>,
  90. ) -> Result<i32> {
  91. let file = Arc::new(Mutex::new(if append_log {
  92. fs::OpenOptions::new()
  93. .create(true)
  94. .append(true)
  95. .open(log_path)?
  96. } else {
  97. fs::File::create(log_path)?
  98. }));
  99. // The command: local (the TTY wrapper or direct exec) or remote (SSH to
  100. // the host — a pty via `ssh -t` keeps the rich Live display from
  101. // deadlocking on the non-tty pipe; `stty -onlcr` keeps the log LF-only;
  102. // TTY_COMPATIBLE/NO_COLOR match the local run). For the remote host the
  103. // `cwd` is a remote path embedded in the shell string, so the local
  104. // current_dir/env are skipped.
  105. let mut cmd = if let Some(host) = ssh_host {
  106. crate::bench_remote::remote_command(host, cwd, bin, args)
  107. } else if cfg!(target_os = "macos")
  108. && std::process::Command::new("/usr/bin/script")
  109. .arg("-q")
  110. .arg("/dev/null")
  111. .arg("true")
  112. .output()
  113. .is_ok()
  114. {
  115. // Wrap in `sh` + `stty -onlcr`: the macOS pty's ONLCR would turn every
  116. // `\n` the tool writes into `\r\n` on disk, and the report viewer
  117. // then renders each line twice (blank line after every line).
  118. // `stty -onlcr` disables the post-processing at the source; `exec`
  119. // replaces the shell so the tool keeps the pty and the exit code.
  120. let mut c = std::process::Command::new("/usr/bin/script");
  121. c.arg("-q");
  122. c.arg("/dev/null");
  123. c.arg("sh");
  124. c.arg("-c");
  125. c.arg("stty -onlcr 2>/dev/null; exec \"$@\"");
  126. c.arg("sh"); // $0 inside the -c script
  127. c.arg(bin);
  128. c.args(args);
  129. c
  130. } else {
  131. let mut c = std::process::Command::new(bin);
  132. c.args(args);
  133. c
  134. };
  135. if ssh_host.is_none() {
  136. cmd.current_dir(cwd);
  137. cmd.env("TTY_COMPATIBLE", "0");
  138. cmd.env("NO_COLOR", "1");
  139. }
  140. cmd.stdout(Stdio::piped());
  141. cmd.stderr(Stdio::piped());
  142. let mut child = cmd.spawn()?;
  143. let stdout_pipe = child.stdout.take().expect("stdout pipe");
  144. let stderr_pipe = child.stderr.take().expect("stderr pipe");
  145. let f1 = file.clone();
  146. let t1 = std::thread::spawn(move || crate::bench_io::pump(stdout_pipe, true, f1));
  147. let t2 = std::thread::spawn(move || crate::bench_io::pump(stderr_pipe, false, file));
  148. let status = child.wait()?;
  149. t1.join().ok();
  150. t2.join().ok();
  151. Ok(status.code().unwrap_or(-1))
  152. }