launcher/src/exec.rs

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

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

/// 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!(fluent_i18n::t!("exec-both-specified"))
        }
        (None, None) => {
            anyhow::bail!(fluent_i18n::t!("exec-missing"))
        }
        (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, &config.recipe_sources).await?;
            let cname = recipe.cluster.as_deref().unwrap_or("auto");
            let (cluster, _) =
                select_cluster(config, cname, recipe.required_nodes, &HashSet::new())
                    .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 {
            let name = n.to_string();
            anyhow::bail!(
                fluent_i18n::t!("exec-node-not-found", { "name" => &name, "cluster" => &cluster.name })
            );
        }
        anyhow::bail!(fluent_i18n::t!("exec-cluster-empty", { "name" => &cluster.name }));
    }

    let cmd = command.to_string();
    info!(
        "{}",
        fluent_i18n::t!("exec-header", { "cmd" => &cmd, "nodes" => 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(crate::locale::spawn_with_locale(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!("{}", fluent_i18n::t!("exec-section", { "name" => &name }));
                    print!("{}", output);
                    if !output.ends_with('\n') {
                        println!();
                    }
                }
            }
            Err(e) => {
                any_failed = true;
                warn!(
                    "{}",
                    fluent_i18n::t!("exec-failed", { "name" => &name, "cmd" => &command, "error" => &e.to_string() })
                );
            }
        }
    }

    if any_failed {
        warn!("{}", fluent_i18n::t!("exec-completed-errors"));
    }
    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").to_string();
            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!(fluent_i18n::t!("exec-unknown-alias-empty", { "name" => &name }))
                } else {
                    anyhow!(fluent_i18n::t!("exec-unknown-alias-list", { "name" => &name, "aliases" => &known.join(", ") }))
                }
            })
        }
        (true, true) => {
            anyhow::bail!(fluent_i18n::t!("exec-cannot-use-alias"))
        }
        (false, false) => {
            anyhow::bail!(fluent_i18n::t!("exec-must-specify"))
        }
        (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".into(),
            clusters: vec![],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases,
            home_assistant: None,
            mikrotik: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            recipe_sources: Vec::new(),
            docker_image_sources: Vec::new(),
            cluster_sources: Vec::new(),
            global_mods_dir: None,
        }
    }

    #[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");
    }
}