Przeglądaj źródła

bench: add sm121 bench — tool-eval-bench with auto-landed reports; move stop handler to docker/stop

Alexander Korolev 1 miesiąc temu
rodzic
commit
1ef3eebed6
6 zmienionych plików z 426 dodań i 48 usunięć
  1. 7 1
      AGENTS.md
  2. 340 0
      src/bench.rs
  3. 15 0
      src/cli.rs
  4. 44 0
      src/docker/stop.rs
  5. 19 46
      src/main.rs
  6. 1 1
      src/web/api.rs

+ 7 - 1
AGENTS.md

@@ -17,13 +17,19 @@ See [`DEVELOPMENT.md`](DEVELOPMENT.md) for architecture, module layout, and dev
 ## Commands
 
 ```
-serve/stop/logs/health/doctor/exec/power/tmux/init-host-keys/config/web
+serve/stop/bench/logs/health/doctor/exec/power/tmux/init-host-keys/config/web
 ```
 
 `stop` waits until the containers are actually stopped and removed (polls up
 60 s per node) — an immediate `serve` on the same nodes would otherwise race
 the dying containers' GPU release. `--no-wait` opts out.
 
+`bench` runs `tool-eval-bench` (from PATH) against a live serve and lands
+`<UTC-date>_<label>.{md,json,log}` in `recipes/<lane>/reports/` — it does the
+scratch-cwd/absolute-json/log juggling, polls the API for readiness first,
+and prints the final score + perf rows. Defaults: seed 42, hardmode, perf,
+concurrency 1,2; served name/port from the recipe.
+
 `sm121 web` — local axum + SSE dashboard: LED board (cluster-level health,
 non-blocking `extra_entities` as ⚙ badges), themes, serve/stop/power jobs
 with type-the-name power gate, live log-less jobs window fed by docker

+ 340 - 0
src/bench.rs

@@ -0,0 +1,340 @@
+//! `sm121 bench` — run `tool-eval-bench` against a live serve and land the
+//! reports in the repo.
+//!
+//! The bench tool needs babysitting this command does instead of the
+//! operator: it must run in a scratch cwd (it writes
+//! `runs/<date>/<run>.md` relative to cwd), needs an absolute `--json-file`
+//! path, and the `.md` report must be copied into
+//! `recipes/<lane>/reports/<UTC-date>_<label>.{md,json,log}` by hand.
+//! Campaign convention (seed 42, hardmode, perf, concurrency 1,2) is the
+//! default; `sm121 bench -d` prints the exact command.
+
+use std::fs;
+use std::io::{Read, Write};
+use std::path::{Path, PathBuf};
+use std::process::{Command, Stdio};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
+
+use anyhow::Result;
+
+use crate::config::{
+    Config, Recipe, default_config_dir, load_recipe, local_path, select_cluster_by_name,
+};
+
+/// Fields of the `sm121 bench` subcommand (mirrored by the clap variant in
+/// `cli.rs`). Served model name and port come from the recipe
+/// (`--served-model-name` in args / `port`); the bench binary is looked up on PATH.
+#[derive(Clone)]
+pub struct Args {
+    pub recipe: String,
+    pub cluster: String,
+    pub label: String,
+    pub concurrency: String,
+    pub seed: u64,
+}
+
+/// How long to poll the API for readiness before giving up.
+const READY_TIMEOUT_SECS: u64 = 900;
+/// Readiness poll interval.
+const READY_POLL_SECS: u64 = 5;
+
+/// Locate the recipe file (same search order as [`load_recipe`]) so reports
+/// land next to it.
+fn recipe_file(recipe_name: &str, sources: &[String]) -> Option<PathBuf> {
+    let mut paths: Vec<PathBuf> = vec![
+        PathBuf::from(format!("recipes/{}/recipe.toml", recipe_name)),
+        PathBuf::from(format!("{}/recipe.toml", recipe_name)),
+    ];
+    for src in sources {
+        paths.push(local_path(src).join(recipe_name).join("recipe.toml"));
+    }
+    if let Some(dir) = default_config_dir() {
+        paths.push(dir.join("recipes").join(recipe_name).join("recipe.toml"));
+    }
+    paths.into_iter().find(|p| p.is_file())
+}
+
+/// Served model name from the recipe's `--served-model-name` arg.
+fn served_model(recipe: &Recipe) -> Option<String> {
+    let args = recipe.args.to_vec();
+    args.iter()
+        .position(|s| s == "--served-model-name")
+        .and_then(|i| args.get(i + 1))
+        .cloned()
+}
+
+/// `YYYY-MM-DD_HHMMSS` in UTC (no chrono dependency).
+fn utc_stamp() -> String {
+    let secs: i64 = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_secs() as i64)
+        .unwrap_or(0);
+    let (h, m, s) = ((secs % 86400) / 3600, (secs % 3600) / 60, secs % 60);
+    let (y, mo, d) = civil_date(secs / 86400);
+    format!(
+        "{y:04}-{mo:02}-{d:02}_{h:02}{m:02}{s:02}",
+        y = y,
+        mo = mo,
+        d = d
+    )
+}
+
+/// Days since 1970-01-01 → (year, month, day); Howard Hinnant's civil_from_days.
+fn civil_date(z: i64) -> (i64, i64, i64) {
+    let z = z + 719_468;
+    let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
+    let doe = z - era * 146_097;
+    let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365;
+    let y = yoe + era * 400;
+    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
+    let mp = (5 * doy + 2) / 153;
+    let d = doy - (153 * mp + 2) / 5 + 1;
+    let mo = if mp < 10 { mp + 3 } else { mp - 9 };
+    let y = if mo <= 2 { y + 1 } else { y };
+    (y, mo, d)
+}
+
+/// Keep the label filesystem-safe.
+fn sanitize(label: &str) -> String {
+    label
+        .chars()
+        .map(|c| {
+            if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
+                c
+            } else {
+                '-'
+            }
+        })
+        .collect()
+}
+
+/// Poll `GET <base>/models` until it answers (engine up).
+async fn wait_ready(base_url: &str) -> Result<()> {
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(10))
+        .build()?;
+    let deadline = Instant::now() + Duration::from_secs(READY_TIMEOUT_SECS);
+    let mut since_last_note = Instant::now();
+    loop {
+        match client.get(format!("{base_url}/models")).send().await {
+            Ok(r) if r.status().is_success() => return Ok(()),
+            Ok(_) => {}
+            Err(_) => {}
+        }
+        if since_last_note.elapsed() >= Duration::from_secs(30) {
+            since_last_note = Instant::now();
+            eprintln!(
+                "bench: waiting for API at {base_url} (up to {} s)...",
+                READY_TIMEOUT_SECS
+            );
+        }
+        if Instant::now() > deadline {
+            anyhow::bail!("bench: API at {base_url} not ready after {READY_TIMEOUT_SECS} s");
+        }
+        tokio::time::sleep(Duration::from_secs(READY_POLL_SECS)).await;
+    }
+}
+
+/// Tee one output pipe to the console and the log file.
+fn pump<R: Read>(mut reader: R, to_console: bool, file: Arc<Mutex<fs::File>>) {
+    let mut buf = Vec::new();
+    loop {
+        let n = reader.read(&mut buf).expect("reading bench output pipe");
+        if n == 0 {
+            break;
+        }
+        if to_console {
+            std::io::stdout().write_all(&buf[..n]).ok();
+        }
+        if let Ok(mut f) = file.lock() {
+            f.write_all(&buf[..n]).ok();
+        }
+        buf.clear();
+    }
+}
+
+/// Run the bench tool in `cwd`, teeing stdout+stderr to the console and the
+/// log file; returns the exit code.
+fn run_tool(bin: &str, args: &[String], cwd: &Path, log_path: &Path) -> Result<i32> {
+    let file = Arc::new(Mutex::new(fs::File::create(log_path)?));
+    let mut child = Command::new(bin)
+        .args(args)
+        .current_dir(cwd)
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .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 || pump(stdout_pipe, true, f1));
+    let t2 = std::thread::spawn(move || pump(stderr_pipe, false, file));
+    let status = child.wait()?;
+    t1.join().ok();
+    t2.join().ok();
+    Ok(status.code().unwrap_or(-1))
+}
+
+/// Newest `.md` under `scratch/runs/` (the tool's `runs/<date>/<run>.md`).
+fn latest_md(scratch: &Path) -> Option<PathBuf> {
+    let runs = scratch.join("runs");
+    let mut best: Option<(std::time::SystemTime, PathBuf)> = None;
+    let mut stack = vec![runs];
+    while let Some(dir) = stack.pop() {
+        let entries = fs::read_dir(dir).ok()?;
+        for e in entries.flatten() {
+            let p = e.path();
+            if p.is_dir() {
+                stack.push(p);
+            } else if p.extension().is_some_and(|x| x == "md") {
+                let t = e
+                    .metadata()
+                    .and_then(|m| m.modified())
+                    .ok()
+                    .unwrap_or(SystemTime::UNIX_EPOCH);
+                if best.as_ref().is_none_or(|(bt, _)| t > *bt) {
+                    best = Some((t, p));
+                }
+            }
+        }
+    }
+    best.map(|(_, p)| p)
+}
+
+/// 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 = 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 = match cluster.nodes.iter().find(|n| n.rank == 0) {
+        Some(n) => n,
+        None => cluster
+            .nodes
+            .first()
+            .ok_or_else(|| anyhow::anyhow!("cluster '{}' has no nodes", cluster.name))?,
+    };
+    let port = 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 = utc_stamp();
+    let label = sanitize(&a.label);
+    let reports_dir = file
+        .parent()
+        .ok_or_else(|| anyhow::anyhow!("no parent dir for {}", file.display()))?
+        .join("reports");
+    let json_path = reports_dir.join(format!("{stamp}_{label}.json"));
+    let log_path = reports_dir.join(format!("{stamp}_{label}.log"));
+    let scratch = dirs::home_dir()
+        .ok_or_else(|| anyhow::anyhow!("no home dir"))?
+        .join(".cache")
+        .join("sm121")
+        .join("bench")
+        .join(format!("{stamp}_{label}"));
+    let bin = "tool-eval-bench".to_string();
+
+    let args = vec![
+        "--base-url".into(),
+        base_url.clone(),
+        "--model".into(),
+        model.clone(),
+        "--seed".into(),
+        a.seed.to_string(),
+        "--hardmode".into(),
+        "--perf".into(),
+        "--concurrency".into(),
+        a.concurrency.clone(),
+        "--json-file".into(),
+        json_path.to_string_lossy().to_string(),
+    ];
+
+    if dry_run {
+        println!("bench: cwd {}", scratch.display());
+        println!("bench: {} {}", bin, args.join(" "));
+        println!("bench: reports -> {}", reports_dir.display());
+        return Ok(());
+    }
+
+    fs::create_dir_all(&scratch)?;
+    fs::create_dir_all(&reports_dir)?;
+
+    eprintln!("bench: waiting for {base_url} ...");
+    wait_ready(&base_url).await?;
+    eprintln!("bench: API ready — running {} (seed {})", bin, a.seed);
+
+    let code = run_tool(&bin, &args, &scratch, &log_path)?;
+    if code != 0 {
+        anyhow::bail!(
+            "bench: tool exited with {code} (log: {})",
+            log_path.display()
+        );
+    }
+
+    let md_src = 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)?;
+
+    // Summary: final score from the JSON + the perf table from the log.
+    let j: serde_json::Value = serde_json::from_str(&fs::read_to_string(&json_path)?)?;
+    let sc = &j["scores"];
+    let completion = sc["completion_rate"]
+        .as_f64()
+        .map(|v| v.to_string())
+        .unwrap_or_default();
+    eprintln!(
+        "bench: final score {} / {} ({} points), completion {completion}",
+        sc["final_score"], sc["max_points"], sc["total_points"]
+    );
+    for line in fs::read_to_string(&log_path)?.lines() {
+        if line.contains("pp2048") && line.contains("tg128") {
+            println!("{line}");
+        }
+    }
+    eprintln!("bench: reports landed in {}", reports_dir.display());
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn stamp_format() {
+        let s = utc_stamp();
+        assert_eq!(s.len(), 17, "{s}");
+        assert_eq!(&s[4..5], "-");
+        assert_eq!(&s[7..8], "-");
+        assert_eq!(&s[10..11], "_");
+    }
+
+    #[test]
+    fn civil_date_known_values() {
+        assert_eq!(civil_date(0), (1970, 1, 1));
+        assert_eq!(civil_date(19_723), (2024, 1, 1));
+        assert_eq!(civil_date(20_818), (2026, 12, 31));
+    }
+
+    #[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");
+    }
+
+    #[test]
+    fn served_model_parsed_from_args() {
+        let content = fs::read_to_string("recipes/glm53-flash-r27-spark/recipe.toml")
+            .expect("repo recipe for test");
+        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());
+    }
+}

+ 15 - 0
src/cli.rs

@@ -144,6 +144,21 @@ pub enum Commands {
         no_wait: bool,
     },
 
+    /// Bench a live serve with tool-eval-bench: reports auto-land in
+    /// recipes/<lane>/reports/<UTC-date>_<label>.{md,json,log} (seed 42, hardmode, perf).
+    Bench {
+        #[arg(long, required = true)]
+        recipe: String,
+        #[arg(long, required = true)]
+        cluster: String,
+        #[arg(long, required = true)]
+        label: String,
+        #[arg(long, default_value = "1,2")]
+        concurrency: String,
+        #[arg(long, default_value_t = 42)]
+        seed: u64,
+    },
+
     /// Check cluster health
     Health {
         #[arg(

+ 44 - 0
src/docker/stop.rs

@@ -43,6 +43,50 @@ pub fn remote_command(recipe_name: &str, cluster_name: &str, rank: u32, wait: bo
     )
 }
 
+/// Handle the `stop` command: stop containers and optionally power off.
+///
+/// Waits for the containers to actually stop (see [`remote_command`]), then
+/// runs the recipe's `after_stop` hooks and the optional power off.
+pub async fn handle(
+    config: &crate::config::Config,
+    recipe_name: &str,
+    cluster_name: &str,
+    nodes: &[String],
+    wait: bool,
+    power_off: bool,
+    dry_run: bool,
+) -> Result<(), anyhow::Error> {
+    let (cluster, nodes): (crate::config::Cluster, Vec<String>) = if !nodes.is_empty() {
+        let cluster = crate::resolve_cluster_for_op(config, recipe_name, cluster_name).await?;
+        (cluster, nodes.to_vec())
+    } else {
+        let variants =
+            crate::docker::deploy::discover_recipe_deployments(config, recipe_name).await?;
+        if variants.is_empty() {
+            anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
+        }
+        if variants.len() == 1 {
+            let v = &variants[0];
+            let cluster = crate::config::select_cluster_by_name(config, &v.cluster)?;
+            (cluster, v.nodes.clone())
+        } else {
+            anyhow::bail!(crate::docker::deploy::multi_deployment_error(
+                recipe_name,
+                &variants
+            ));
+        }
+    };
+    crate::docker::stop_cluster(config, &cluster, recipe_name, &nodes, wait, dry_run).await?;
+    // Run recipe `after_stop` hooks (e.g. restore_dgx_services) on the stopped nodes; best-effort.
+    if let Ok(recipe) = crate::config::load_recipe(recipe_name, &config.recipe_sources).await {
+        crate::after_stop::run(config, &cluster, &recipe, dry_run).await?;
+    }
+    if power_off {
+        crate::power::power_off(config, &cluster, &nodes, dry_run).await?;
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::remote_command;

+ 19 - 46
src/main.rs

@@ -7,6 +7,7 @@ fluent_i18n::i18n!("locales", fallback = "en");
 
 mod after_stop;
 mod before_serve;
+mod bench;
 mod builtin_scripts;
 mod cli;
 mod cluster_state;
@@ -142,7 +143,7 @@ async fn main() -> Result<()> {
             power_off,
             no_wait,
         } => {
-            handle_stop(
+            crate::docker::stop::handle(
                 &config,
                 &recipe_name,
                 &cluster,
@@ -153,6 +154,22 @@ async fn main() -> Result<()> {
             )
             .await?;
         }
+        Commands::Bench {
+            recipe,
+            cluster,
+            label,
+            concurrency,
+            seed,
+        } => {
+            let args = crate::bench::Args {
+                recipe,
+                cluster,
+                label,
+                concurrency,
+                seed,
+            };
+            crate::bench::run(&config, &args, dry_run).await?;
+        }
         Commands::Health {
             cluster: cluster_name,
             fingerprint,
@@ -335,50 +352,6 @@ async fn handle_logs(
     .await?;
     Ok(())
 }
-/// Handle `stop` command: stop containers and optionally power off.
-pub(crate) async fn handle_stop(
-    config: &Config,
-    recipe_name: &str,
-    cluster: &str,
-    node: &[String],
-    wait: bool,
-    power_off: bool,
-    dry_run: bool,
-) -> Result<()> {
-    let (cluster, nodes): (Cluster, Vec<String>) = if !node.is_empty() {
-        let c = resolve_cluster_for_op(config, recipe_name, cluster).await?;
-        (c, node.to_vec())
-    } else {
-        let variants = docker::deploy::discover_recipe_deployments(config, recipe_name).await?;
-        if variants.is_empty() {
-            anyhow::bail!("No running containers found for recipe '{}'", recipe_name);
-        }
-        if variants.len() == 1 {
-            let v = &variants[0];
-            let c = select_cluster_by_name(config, &v.cluster)?;
-            (c, v.nodes.clone())
-        } else {
-            anyhow::bail!(docker::deploy::multi_deployment_error(
-                recipe_name,
-                &variants
-            ));
-        }
-    };
-    docker::stop_cluster(config, &cluster, recipe_name, &nodes, wait, dry_run).await?;
-
-    // Run recipe `after_stop` hooks (e.g. restore_dgx_services) on the
-    // nodes that were actually stopped. Best-effort against the loaded
-    // recipe; if the recipe file is unavailable, skip silently.
-    if let Ok(recipe) = crate::config::load_recipe(recipe_name, &config.recipe_sources).await {
-        crate::after_stop::run(config, &cluster, &recipe, dry_run).await?;
-    }
-
-    if power_off {
-        crate::power::power_off(config, &cluster, &nodes, dry_run).await?;
-    }
-    Ok(())
-}
-
 /// Handle `health` command: check cluster health or show host key fingerprints.
 async fn handle_health(
     config: &Config,
@@ -466,7 +439,7 @@ async fn handle_power(config: &Config, action: PowerAction, dry_run: bool) -> Re
 /// Resolve a cluster for container operations (stop/logs).
 /// When `auto`, tries live container label lookup first, then falls back
 /// to config-based auto-selection (for dry-run / offline scenarios).
-async fn resolve_cluster_for_op(
+pub(crate) async fn resolve_cluster_for_op(
     config: &Config,
     recipe_name: &str,
     cluster_name: &str,

+ 1 - 1
src/web/api.rs

@@ -107,7 +107,7 @@ pub(crate) async fn stop(
     super::jobs::spawn(app2, id.clone(), async move {
         let nodes = vec![node.unwrap_or_default()];
         let nodes = if nodes[0].is_empty() { vec![] } else { nodes };
-        crate::handle_stop(&config, &recipe, &cluster, &nodes, true, false, false).await
+        crate::docker::stop::handle(&config, &recipe, &cluster, &nodes, true, false, false).await
     });
     Ok(Json(serde_json::json!({ "job": id })))
 }