Эх сурвалжийг харах

bench: add --node <name> — pin the bench target within a cluster (default: rank 0). A serve launched with 'serve --node' lives on a non-rank-0 node (solo Bolt on dgx3 inside 4nodes-switch), and benching the rank-0 head measured a different engine (the live kindling on dgx1 — model_not_available). pick_head tolerates the .lan suffix; web API passes node: None.

Alexander Korolev 6 өдөр өмнө
parent
commit
8a683153dd
4 өөрчлөгдсөн 68 нэмэгдсэн , 6 устгасан
  1. 60 6
      src/bench.rs
  2. 5 0
      src/cli.rs
  3. 2 0
      src/main.rs
  4. 1 0
      src/web/api.rs

+ 60 - 6
src/bench.rs

@@ -40,6 +40,11 @@ pub struct Args {
     /// is uploaded to the host; the report + run package land locally. The
     /// host must have `tool-eval-bench` on PATH and reach the serve.
     pub ssh_host: Option<String>,
+    /// Bench target node within the cluster (default: rank 0). Use when the
+    /// serve was launched with `serve --node <n>` — a pinned/solo run may
+    /// live on a non-rank-0 node, and benching the wrong head measures a
+    /// different (or no) engine.
+    pub node: Option<String>,
     /// Disable thinking/reasoning in the harness (tool `--no-think`, sets
     /// `enable_thinking=false`). Needed for thinking-by-default models (Ling,
     /// GLM-4.5-style) whose long default reasoning blows the per-turn timeout
@@ -59,6 +64,34 @@ fn served_model(recipe: &Recipe) -> Option<String> {
         .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).
+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.
 fn sanitize(label: &str) -> String {
     label
@@ -170,12 +203,7 @@ pub async fn run(config: &Config, a: &Args, dry_run: bool) -> Result<()> {
         .ok_or_else(|| anyhow::anyhow!("recipe file not found for '{}'", a.recipe))?
         .canonicalize()?;
     let cluster = select_cluster_by_name(config, &a.cluster)?;
-    let head = cluster
-        .nodes
-        .iter()
-        .find(|n| n.rank == 0)
-        .or_else(|| cluster.nodes.first())
-        .ok_or_else(|| anyhow::anyhow!("cluster '{}' has no nodes", cluster.name))?;
+    let head = pick_head(&cluster.nodes, &cluster.name, a.node.as_deref())?;
     let port = recipe.port;
     let model = served_model(&recipe).ok_or_else(|| {
         anyhow::anyhow!(
@@ -411,8 +439,34 @@ pub async fn run(config: &Config, a: &Args, dry_run: bool) -> Result<()> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::config::Node;
     use std::process::Command;
 
+    fn node(name: &str, rank: u32) -> Node {
+        Node {
+            name: name.to_string(),
+            ssh_ip: "192.0.2.1".into(),
+            ib_ip: None,
+            rank,
+            host_key: None,
+            entity: None,
+        }
+    }
+
+    #[test]
+    fn pick_head_defaults_to_rank0() {
+        let nodes = vec![node("a.lan", 2), node("b.lan", 0), node("c.lan", 1)];
+        assert_eq!(pick_head(&nodes, "c", None).unwrap().name, "b.lan");
+    }
+
+    #[test]
+    fn pick_head_node_override_tolerates_lan_suffix() {
+        let nodes = vec![node("a.lan", 2), node("b.lan", 0)];
+        assert_eq!(pick_head(&nodes, "c", Some("a")).unwrap().name, "a.lan");
+        assert_eq!(pick_head(&nodes, "c", Some("b.lan")).unwrap().name, "b.lan");
+        assert!(pick_head(&nodes, "c", Some("zzz")).is_err());
+    }
+
     /// Regression: the pump must not treat a not-yet-written pipe as EOF.
     /// `read()` into a zero-length buffer returns 0, which the old pump
     /// treated as EOF — closing the pipe and killing the child on its first

+ 5 - 0
src/cli.rs

@@ -183,6 +183,11 @@ pub enum Commands {
         recipe: String,
         #[arg(long, required = true)]
         cluster: String,
+        /// Bench target node within the cluster (default: rank 0). Use when
+        /// the serve was launched with `serve --node <n>` — a pinned/solo
+        /// run may live on a non-rank-0 node.
+        #[arg(long)]
+        node: Option<String>,
         /// Label for the run (campaign record). Optional for `--perf-only`
         /// runs (default `perf`); required for quality runs.
         #[arg(long)]

+ 2 - 0
src/main.rs

@@ -251,6 +251,7 @@ async fn main() -> Result<()> {
         Commands::Bench {
             recipe,
             cluster,
+            node,
             label,
             concurrency,
             seed,
@@ -265,6 +266,7 @@ async fn main() -> Result<()> {
             let args = crate::bench::Args {
                 recipe,
                 cluster,
+                node,
                 label,
                 concurrency,
                 seed,

+ 1 - 0
src/web/api.rs

@@ -156,6 +156,7 @@ pub(crate) async fn bench(
     let args = crate::bench::Args {
         recipe: req.recipe.clone(),
         cluster: req.cluster.clone(),
+        node: None,
         label,
         concurrency: String::from("1,2"),
         seed: 42,