bench_remote.rs 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. //! Remote-host bench execution (`sm121 bench --ssh-host <host>`): the
  2. //! tool-eval-bench runs on the remote host over SSH (the local machine only
  3. //! orchestrates), keeping the bench client's own network (a WiFi laptop) off
  4. //! the measurement path. The tokenizer is uploaded to the host; the report +
  5. //! run package still land locally (the remote json/md are downloaded back).
  6. use std::path::Path;
  7. use std::process::Command;
  8. use anyhow::Result;
  9. /// Build the SSH command that runs the tool on the remote host. `ssh -t`
  10. /// allocates a pty (the rich Live display deadlocks on a non-tty pipe);
  11. /// `stty -onlcr` keeps the log LF-only; TTY_COMPATIBLE/NO_COLOR match the
  12. /// local run. The `cwd` is a remote path embedded in the shell string.
  13. pub fn remote_command(host: &str, cwd: &Path, bin: &str, args: &[String]) -> Command {
  14. let remote = format!(
  15. "mkdir -p '{}' && cd '{}' || exit 127; stty -onlcr 2>/dev/null; TTY_COMPATIBLE=0 NO_COLOR=1 exec {} {}",
  16. cwd.display(),
  17. cwd.display(),
  18. bin,
  19. args.join(" ")
  20. );
  21. let mut c = Command::new("ssh");
  22. c.arg("-t");
  23. c.arg(host);
  24. c.arg(&remote);
  25. c
  26. }
  27. /// Create the remote scratch dir + upload the tokenizer there so the bench
  28. /// client's own network stays off the measurement path.
  29. pub fn prepare(host: &str, tool_scratch: &Path, tokenizer: &Option<String>) -> Result<()> {
  30. let st = Command::new("ssh")
  31. .arg(host)
  32. .arg(format!("mkdir -p '{}'", tool_scratch.display()))
  33. .output()
  34. .map_err(|e| anyhow::anyhow!("ssh mkdir to {host}: {e}"))?;
  35. if !st.status.success() {
  36. anyhow::bail!("mkdir on {host}: {}", String::from_utf8_lossy(&st.stderr));
  37. }
  38. if let Some(tok) = tokenizer {
  39. let st = Command::new("scp")
  40. .arg("-r")
  41. .arg(tok)
  42. .arg(format!("{host}:{}", tool_scratch.display()))
  43. .output()
  44. .map_err(|e| anyhow::anyhow!("scp tokenizer to {host}: {e}"))?;
  45. if !st.status.success() {
  46. anyhow::bail!(
  47. "scp tokenizer to {host}: {}",
  48. String::from_utf8_lossy(&st.stderr)
  49. );
  50. }
  51. eprintln!("bench: tokenizer uploaded to {host}");
  52. }
  53. Ok(())
  54. }
  55. /// Download the json + the md from the remote scratch dir so the local report
  56. /// generation is unchanged. A download failure warns (the run is not lost);
  57. /// the local json is simply missing (the run package is skipped). Perf-only
  58. /// runs write no JSON by design — pass `expect_json = false` to skip the
  59. /// scp (and the warning) entirely.
  60. pub fn collect(
  61. host: &str,
  62. tool_json: &Path,
  63. json_path: &Path,
  64. tool_scratch: &Path,
  65. scratch: &Path,
  66. expect_json: bool,
  67. ) -> Result<()> {
  68. if expect_json {
  69. let st = Command::new("scp")
  70. .arg(format!("{host}:{}", tool_json.display()))
  71. .arg(json_path)
  72. .output()
  73. .map_err(|e| anyhow::anyhow!("scp json from {host}: {e}"))?;
  74. if !st.status.success() {
  75. eprintln!("bench: WARNING could not download json from {host}");
  76. }
  77. }
  78. let st = Command::new("ssh")
  79. .arg(host)
  80. .arg(format!(
  81. "find '{}' -name '*.md' -printf '%T@\\t%p\\n' 2>/dev/null | sort -n | tail -1 | cut -f2",
  82. tool_scratch.join("runs").display()
  83. ))
  84. .output()
  85. .map_err(|e| anyhow::anyhow!("ssh find md on {host}: {e}"))?;
  86. let remote_md = String::from_utf8_lossy(&st.stdout).trim().to_string();
  87. if !remote_md.is_empty() {
  88. let md_name = remote_md
  89. .rsplit('/')
  90. .next()
  91. .unwrap_or("report.md")
  92. .to_string();
  93. let local_runs = scratch.join("runs");
  94. std::fs::create_dir_all(&local_runs)?;
  95. let local_md_path = local_runs.join(&md_name);
  96. let st = Command::new("scp")
  97. .arg(format!("{host}:{remote_md}"))
  98. .arg(&local_md_path)
  99. .output()
  100. .map_err(|e| anyhow::anyhow!("scp md from {host}: {e}"))?;
  101. if !st.status.success() {
  102. eprintln!("bench: WARNING could not download md from {host}");
  103. }
  104. }
  105. Ok(())
  106. }