launcher/src/exec.rs

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

use crate::config::{Cluster, Config, Recipe, select_cluster};

/// Resolve the target cluster for `cargo run exec` from `--cluster` or
/// `--recipe` (exactly one of which must be given).
pub async fn resolve_target(
    config: &Config,
    cluster_name: Option<&str>,
    recipe_name: Option<&str>,
) -> Result<(Cluster, Option<Recipe>)> {
    match (cluster_name, recipe_name) {
        (Some(_), Some(_)) => {
            anyhow::bail!("Specify either --cluster or --recipe, not both")
        }
        (None, None) => {
            anyhow::bail!("Must specify either --cluster or --recipe")
        }
        (Some(name), None) => {
            let cluster = crate::config::select_cluster_by_name(config, name)?;
            Ok((cluster, None))
        }
        (None, Some(rname)) => {
            let recipe = crate::config::load_recipe(rname).await?;
            let cname = recipe.cluster.as_deref().unwrap_or("auto");
            let (cluster, _) = select_cluster(config, cname)
                .with_context(|| format!("resolving cluster for recipe '{}'", rname))?;
            Ok((cluster, Some(recipe)))
        }
    }
}

/// Run `command` on every node of `cluster` (or a single node if `node`
/// is given) and stream the output.  Used by the `cargo run exec`
/// subcommand to support mass-execution flows like `sudo shutdown now`.
pub async fn run(
    config: &Config,
    cluster: &Cluster,
    node_name: Option<&str>,
    command: &str,
    dry_run: bool,
) -> Result<()> {
    let targets: Vec<_> = match node_name {
        Some(name) => cluster
            .nodes
            .iter()
            .filter(|n| n.name == name)
            .cloned()
            .collect(),
        None => cluster.nodes.clone(),
    };

    if targets.is_empty() {
        if let Some(n) = node_name {
            anyhow::bail!("Node '{}' not found in cluster '{}'", n, cluster.name);
        }
        anyhow::bail!("Cluster '{}' has no nodes", cluster.name);
    }

    info!("⚡ exec: {} ({} node(s))", command, targets.len());

    let mut handles = Vec::with_capacity(targets.len());
    for node in targets {
        let config = config.clone();
        let command = command.to_string();
        let dry = dry_run;
        handles.push(tokio::spawn(async move {
            let name = node.name.clone();
            let result = crate::ssh::run_ssh(&config, &node, command.clone(), dry).await;
            (name, command, result)
        }));
    }

    let mut any_failed = false;
    for h in handles {
        let (name, command, result) = h.await?;
        match result {
            Ok(output) => {
                if !output.trim().is_empty() {
                    println!("\n--- {} ---", name);
                    print!("{}", output);
                    if !output.ends_with('\n') {
                        println!();
                    }
                }
            }
            Err(e) => {
                any_failed = true;
                warn!("✗ {} :: {} :: {}", name, command, e);
            }
        }
    }

    if any_failed {
        warn!("exec completed with errors (some nodes reported failures)");
    }
    Ok(())
}

/// Resolve the command to execute for `cargo run exec`.
///
/// - `--alias <name>` looks up `name` in `config.exec_aliases` and returns
///   the configured shell command.  Errors if not found.
/// - Otherwise joins the trailing positional args into a single command
///   line (must be non-empty).
///
/// Exactly one of `alias_name` or `command_args` must be provided.
pub fn resolve_command(
    config: &Config,
    alias_name: Option<&str>,
    command_args: &[String],
) -> Result<String> {
    let has_alias = alias_name.is_some();
    let has_args = !command_args.is_empty();
    match (has_alias, has_args) {
        (true, false) => {
            let name = alias_name.expect("checked by has_alias");
            config.exec_aliases.get(name).cloned().ok_or_else(|| {
                let mut known: Vec<String> = config.exec_aliases.keys().cloned().collect();
                known.sort();
                if known.is_empty() {
                    anyhow!("Unknown alias @{} (no aliases configured)", name)
                } else {
                    anyhow!(
                        "Unknown alias @{}; known aliases: {}",
                        name,
                        known.join(", ")
                    )
                }
            })
        }
        (true, true) => {
            anyhow::bail!("Cannot use --alias together with a trailing command")
        }
        (false, false) => {
            anyhow::bail!("Must specify either --alias <name> or a command after `--`")
        }
        (false, true) => Ok(command_args.join(" ")),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn test_config(aliases: &[(&str, &str)]) -> Config {
        let exec_aliases: HashMap<String, String> = aliases
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect();
        Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".to_string(),
            clusters: vec![],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases,
        }
    }

    #[test]
    fn resolve_command_alias() {
        let cfg = test_config(&[("gpu-lock", "sudo nvidia-smi -lgc 300,2300")]);
        let cmd = resolve_command(&cfg, Some("gpu-lock"), &[]).unwrap();
        assert_eq!(cmd, "sudo nvidia-smi -lgc 300,2300");
    }

    #[test]
    fn resolve_command_unknown_alias_lists_known() {
        let cfg = test_config(&[("gpu-lock", "a"), ("shutdown", "b")]);
        let err = resolve_command(&cfg, Some("nope"), &[]).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("Unknown alias @nope"), "got: {msg}");
        assert!(msg.contains("gpu-lock"), "got: {msg}");
        assert!(msg.contains("shutdown"), "got: {msg}");
    }

    #[test]
    fn resolve_command_unknown_alias_no_aliases() {
        let cfg = test_config(&[]);
        let err = resolve_command(&cfg, Some("anything"), &[]).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("no aliases configured"), "got: {msg}");
    }

    #[test]
    fn resolve_command_no_alias_no_args_errors() {
        let cfg = test_config(&[]);
        let err = resolve_command(&cfg, None, &[]).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("Must specify"), "got: {msg}");
    }

    #[test]
    fn resolve_command_alias_and_args_errors() {
        let cfg = test_config(&[("gpu", "a")]);
        let err = resolve_command(&cfg, Some("gpu"), &["ls".into()]).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("Cannot use --alias together"), "got: {msg}");
    }

    #[test]
    fn resolve_command_literal_join() {
        let cfg = test_config(&[]);
        let cmd =
            resolve_command(&cfg, None, &["echo".into(), "hello".into(), "world".into()]).unwrap();
        assert_eq!(cmd, "echo hello world");
    }
}