launcher/src/before_serve.rs

use anyhow::Result;
use tracing::{info, warn};

use crate::config::{Cluster, Config, OnFailure, PreServeCommand, Recipe};

/// Run pre-serve commands on every node of the cluster before launching
/// containers or tmux sessions.  Combines cluster-level config
/// (`[before_serve."<cluster>"]`) with recipe-level overrides
/// (`before_serve = [...]` in recipe.toml).  Cluster-level commands run
/// first, then recipe-level — each in declared order — so the resulting
/// state is predictable.
///
/// Per-command behaviour on a non-zero exit is controlled by
/// [`PreServeCommand::on_failure`]:
/// - `abort`   (default) — fail-fast: abort the whole serve
/// - `warn`    — log warning, continue with remaining commands
/// - `continue`— stay silent, continue
///
/// If `skip` is true the function returns without running anything.
pub async fn run(
    config: &Config,
    cluster: &Cluster,
    recipe: &Recipe,
    skip: bool,
    dry_run: bool,
) -> Result<()> {
    if skip {
        info!("⏭  Skipping pre-serve hooks (--skip-before-serve)");
        return Ok(());
    }

    let cmds = combined_commands(config, cluster, recipe);
    if cmds.is_empty() {
        return Ok(());
    }

    info!(
        "🔧 Pre-serve hooks: {} command(s) on {} node(s)",
        cmds.len(),
        cluster.nodes.len()
    );

    // Run all commands for one node sequentially (so order matters per
    // node), and the nodes in parallel.
    let mut handles = Vec::with_capacity(cluster.nodes.len());
    for node in &cluster.nodes {
        let config = config.clone();
        let node = node.clone();
        let cmds = cmds.clone();
        let dry = dry_run;
        handles.push(tokio::spawn(async move {
            (
                node.name.clone(),
                run_for_node(&config, &node, &cmds, dry).await,
            )
        }));
    }

    let mut any_failed = false;
    for h in handles {
        let (name, result) = h.await?;
        if let Err(e) = result {
            any_failed = true;
            warn!("pre-serve: node {} failed: {}", name, e);
        }
    }

    if any_failed {
        anyhow::bail!("Pre-serve hooks failed on at least one node (see warnings above)");
    }
    Ok(())
}

/// Combine cluster-level + recipe-level commands into one ordered list.
fn combined_commands(
    config: &Config,
    cluster: &Cluster,
    recipe: &Recipe,
) -> Vec<(&'static str, PreServeCommand)> {
    let mut out: Vec<(&'static str, PreServeCommand)> = Vec::new();
    if let Some(table) = config.before_serve.get(&cluster.name) {
        for c in &table.commands {
            out.push(("cluster", c.clone()));
        }
    }
    for c in &recipe.before_serve {
        out.push(("recipe", c.clone()));
    }
    out
}

/// Run a fixed list of commands on a single node, sequentially.
/// Returns the first command that failed under `on_failure = abort`.
async fn run_for_node(
    config: &Config,
    node: &crate::config::Node,
    cmds: &[(&'static str, PreServeCommand)],
    dry_run: bool,
) -> Result<()> {
    for (origin, cmd) in cmds {
        let command = cmd.command();
        let on_failure = cmd.on_failure();
        if dry_run {
            info!("  {} ({}) :: {}", node.name, origin, truncate(command, 120));
            continue;
        }
        let result = crate::ssh::run_ssh(config, node, command.to_string(), dry_run).await;
        let status = match &result {
            Ok(_) => "ok",
            Err(_) => "fail",
        };
        let label = format!("[{}|{}|{}]", node.name, origin, status);
        info!("  {} $ {}", label, command);
        if let Err(ref e) = result {
            match on_failure {
                OnFailure::Abort => {
                    return Err(anyhow::anyhow!(
                        "pre-serve command failed on {}: {}",
                        node.name,
                        e
                    ));
                }
                OnFailure::Warn => {
                    warn!("  ⚠  {} ignored: {}", node.name, e);
                }
                OnFailure::Continue => {
                    // silent
                }
            }
        }
    }
    Ok(())
}

fn truncate(s: &str, max: usize) -> String {
    if s.len() <= max {
        s.to_string()
    } else {
        format!("{}…", &s[..max])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{Node, PreServeCommand};
    use std::collections::HashMap;

    fn test_cluster() -> Cluster {
        Cluster {
            name: "c1".to_string(),
            master_ib_ip: "1.2.3.4".to_string(),
            nodes: vec![
                Node {
                    name: "n0".to_string(),
                    ssh_ip: "10.0.0.1".to_string(),
                    ib_ip: "1.2.3.4".to_string(),
                    rank: 0,
                    host_key: None,
                },
                Node {
                    name: "n1".to_string(),
                    ssh_ip: "10.0.0.2".to_string(),
                    ib_ip: "1.2.3.5".to_string(),
                    rank: 1,
                    host_key: None,
                },
            ],
        }
    }

    fn test_config(before_serve: HashMap<String, crate::config::BeforeServeTable>) -> Config {
        Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".to_string(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve,
            exec_aliases: HashMap::new(),
        }
    }

    #[test]
    fn combined_orders_cluster_then_recipe() {
        let mut bs = HashMap::new();
        bs.insert(
            "c1".to_string(),
            crate::config::BeforeServeTable {
                commands: vec![PreServeCommand::Bare("cluster_cmd_1".into())],
            },
        );
        let cfg = test_config(bs);
        let recipe = Recipe {
            before_serve: vec![
                PreServeCommand::Bare("recipe_cmd_1".into()),
                PreServeCommand::Bare("recipe_cmd_2".into()),
            ],
            ..test_recipe()
        };

        let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
        let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
        let origins: Vec<&str> = cmds.iter().map(|(o, _)| *o).collect();
        assert_eq!(names, vec!["cluster_cmd_1", "recipe_cmd_1", "recipe_cmd_2"]);
        assert_eq!(origins, vec!["cluster", "recipe", "recipe"]);
    }

    #[test]
    fn combined_empty_when_unconfigured() {
        let cfg = test_config(HashMap::new());
        let recipe = Recipe {
            before_serve: vec![],
            ..test_recipe()
        };
        assert!(combined_commands(&cfg, &test_cluster(), &recipe).is_empty());
    }

    fn test_recipe() -> Recipe {
        Recipe {
            engine: "vllm".to_string(),
            name: "Test".to_string(),
            docker_image: String::new(),
            model: "x".to_string(),
            port: 0,
            required_nodes: None,
            pipeline_parallel_size: None,
            tensor_parallel_size: 1,
            args: vec![],
            env: HashMap::new(),
            cluster: None,
            entrypoint: vec![],
            distributed_backend: "mp".to_string(),
            before_serve: vec![],
        }
    }

    // ── Deserialization tests (the ones that live in this module) ──

    #[test]
    fn parses_cluster_level_before_serve() {
        let toml_str = r#"
ssh_user = "u"
ssh_key_path = "/k"

[[clusters]]
name = "c1"
master_ib_ip = "1.1.1.1"
nodes = [
    { name = "n0", ssh_ip = "2.2.2.2", ib_ip = "1.1.1.1", rank = 0 },
]

[before_serve."c1"]
commands = [
    "sudo nvidia-smi -lgc 300,2300",
    { command = "sudo nvidia-smi -pm 1", on_failure = "warn" },
]
"#;
        let cfg: Config = toml::from_str(toml_str).unwrap();
        assert!(cfg.before_serve.contains_key("c1"));
        let list = &cfg.before_serve["c1"].commands;
        assert_eq!(list.len(), 2);
        assert_eq!(list[0].command(), "sudo nvidia-smi -lgc 300,2300");
        assert_eq!(list[0].on_failure(), OnFailure::Abort); // default
        assert_eq!(list[1].command(), "sudo nvidia-smi -pm 1");
        assert_eq!(list[1].on_failure(), OnFailure::Warn);
    }

    #[test]
    fn parses_recipe_level_before_serve() {
        let toml_str = r#"
name = "Demo"
model = "x"
before_serve = [
    "echo a",
    { command = "echo b", on_failure = "continue" },
]
"#;
        let r: Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(r.before_serve.len(), 2);
        assert_eq!(r.before_serve[0].command(), "echo a");
        assert_eq!(r.before_serve[0].on_failure(), OnFailure::Abort);
        assert_eq!(r.before_serve[1].command(), "echo b");
        assert_eq!(r.before_serve[1].on_failure(), OnFailure::Continue);
    }
}