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 config-level hooks with
/// recipe-level overrides (`before_serve = [...]` in recipe.toml).  Config
/// hooks come from `[before_serve."*"]` (global — runs for every cluster)
/// followed by `[before_serve."<cluster>"]` (per-cluster), then recipe —
/// each in declared order — so per-cluster can supplement/override the
/// global list and recipe commands run last.
///
/// 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,
    recipe_name: &str,
    skip: bool,
    dry_run: bool,
) -> Result<()> {
    if skip {
        info!("{}", fluent_i18n::t!("pre-serve-skip"));
        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()
    );

    // A `{ script = "scripts/prepare.sh" }` hook is recipe-relative: resolve
    // the recipe's own directory up front (it lives in the shared recipe repo,
    // next to recipe.toml — not in this binary).
    let recipe_dir = crate::config::recipe_dir(recipe_name, &config.recipe_sources)
        .await
        .ok();

    // 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 (local_rank, node) in cluster.nodes.iter().enumerate() {
        let config = config.clone();
        let node = node.clone();
        let cmds = cmds.clone();
        let dir = recipe_dir.clone();
        let dry = dry_run;
        handles.push(crate::locale::spawn_with_locale(async move {
            (
                node.name.clone(),
                // Serve-LOCAL rank, not `node.rank`: a subset launch (2 of a
                // 4-node cluster) renumbers ranks for the containers, so a
                // per-rank staging hook must see the same 0..N-1 index —
                // otherwise dgx7/dgx8 in a 4-node cluster stage rank 2/3
                // packs for a 2-rank engine (2026-10-21).
                run_for_node(
                    &config,
                    &node,
                    local_rank as u32,
                    &cmds,
                    dir.as_deref(),
                    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!(
                "{}",
                fluent_i18n::t!("pre-serve-failed", { "name" => &name, "error" => e.to_string() })
            );
        }
    }

    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();
    // Global hooks (literal "*" key — not a glob): run for every cluster.
    if let Some(table) = config.before_serve.get("*") {
        for c in &table.commands {
            out.push(("global", c.clone()));
        }
    }
    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,
    local_rank: u32,
    cmds: &[(&'static str, PreServeCommand)],
    recipe_dir: Option<&std::path::Path>,
    dry_run: bool,
) -> Result<()> {
    for (origin, cmd) in cmds {
        let on_failure = cmd.on_failure();
        // Resolve the entry to the SSH command string to run. A `Script`
        // alias is uploaded to the node and invoked via `bash`; a raw
        // command (`Bare`/`Configured`) runs as-is.
        let (label_cmd, ssh_cmd) = match cmd {
            PreServeCommand::Script { script, .. } => {
                // Built-in alias (generic, cluster-wide helpers baked into the
                // binary) or a path relative to the recipe directory (the
                // recipe-specific staging script, shipped in the recipe repo so
                // it evolves with the recipe instead of with a release).
                let (body, tag) = match crate::builtin_scripts::resolve(script) {
                    Some(body) => (body.to_string(), script.clone()),
                    None => (
                        read_recipe_script(script, recipe_dir, &node.name).await?,
                        script.clone(),
                    ),
                };
                let remote = format!(
                    "/tmp/sm121-hook-{}.sh",
                    script
                        .chars()
                        .map(|c| if c.is_ascii_alphanumeric() { c } else { '-' })
                        .collect::<String>()
                );
                crate::ssh::upload_script(config, node, &body, &remote, dry_run).await?;
                // `NODE_RANK` lets a built-in stage per-rank artifacts (e.g. the
                // TensorFold Engram packs) without inventing a second mechanism.
                (
                    format!("script:{tag}"),
                    format!("sudo NODE_RANK={local_rank} bash {remote}"),
                )
            }
            _ => {
                let c = cmd.command();
                (c.to_string(), c.to_string())
            }
        };
        if dry_run {
            info!(
                "{}",
                fluent_i18n::t!("pre-serve-run", { "node" => &node.name, "tag" => &origin, "cmd" => &truncate(&label_cmd, 120) })
            );
            continue;
        }
        let result = crate::ssh::run_ssh(config, node, ssh_cmd, dry_run).await;
        let status = match &result {
            Ok(_) => "ok",
            Err(_) => "fail",
        };
        let label = format!("[{}|{}|{}]", node.name, origin, status);
        info!("  {} $ {}", label, label_cmd);
        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!(
                        "{}",
                        fluent_i18n::t!("pre-serve-ignored", { "cmd" => &node.name, "error" => e.to_string() })
                    );
                }
                OnFailure::Continue => {
                    // silent
                }
            }
        }
    }
    Ok(())
}

fn truncate(s: &str, max: usize) -> String {
    // Iterate by chars (not bytes) so a multibyte boundary inside `max`
    // doesn't panic on `&s[..max]` — pre-serve commands often contain
    // non-ASCII (Cyrillic/CJK paths, en-dashes).
    let head: String = s.chars().take(max).collect();
    if s.chars().count() <= max {
        s.to_string()
    } else {
        format!("{head}…")
    }
}

/// Read a recipe-relative hook script (`{ script = "scripts/prepare.sh" }`).
/// The path must stay inside the recipe directory — no absolute paths, no
/// `..` traversal — so a synced recipe repo cannot make the launcher run an
/// arbitrary file.
async fn read_recipe_script(
    script: &str,
    recipe_dir: Option<&std::path::Path>,
    node: &str,
) -> Result<String> {
    let Some(dir) = recipe_dir else {
        anyhow::bail!(
            "hook script '{script}' is not a built-in alias ({}) and the recipe              directory could not be resolved on node {node}",
            crate::builtin_scripts::known_aliases()
        );
    };
    if script.starts_with('/') || script.split('/').any(|c| c == "..") {
        anyhow::bail!("hook script '{script}' must be a relative path inside the recipe dir");
    }
    let path = dir.join(script);
    tokio::fs::read_to_string(&path).await.map_err(|e| {
        anyhow::anyhow!(
            "hook script '{script}' not found in {} ({e}); it is not a built-in alias either              (known: {})",
            dir.display(),
            crate::builtin_scripts::known_aliases()
        )
    })
}

#[cfg(test)]
#[path = "before_serve_tests.rs"]
mod tests;