launcher/src/after_stop.rs

//! `after_stop` — recipe commands that run on each node **after** containers
//! are stopped (`sm121 stop`). The mirror of [`crate::before_serve`], for the
//! teardown side of a serve: e.g. `restore_dgx_services` to bring back the
//! services that a `before_serve` `stop_dgx_services` stopped.
//!
//! Same entry shape ([`PreServeCommand`]: `Bare` / `Configured` / `Script`)
//! and the same `on_failure` semantics as `before_serve`. Recipe-only (no
//! cluster-level `[after_stop]` config); runs per-node in parallel, commands
//! sequentially within a node.

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

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

/// Run a recipe's `after_stop` commands on every node of the cluster after
/// containers are stopped. See module docs for semantics.
pub async fn run(config: &Config, cluster: &Cluster, recipe: &Recipe, dry_run: bool) -> Result<()> {
    if recipe.after_stop.is_empty() {
        return Ok(());
    }

    info!(
        "🔧 after_stop hooks: {} command(s) on {} node(s)",
        recipe.after_stop.len(),
        cluster.nodes.len()
    );

    let cmds: Vec<(&'static str, PreServeCommand)> = recipe
        .after_stop
        .iter()
        .map(|c| ("recipe", c.clone()))
        .collect();

    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(crate::locale::spawn_with_locale(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!(
                "{}",
                fluent_i18n::t!("after-stop-failed", { "name" => &name, "error" => e.to_string() })
            );
        }
    }

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

/// Run a fixed list of commands on a single node, sequentially. Mirrors
/// `before_serve::run_for_node`; `Script` entries are uploaded + invoked via
/// `bash`, raw entries run as-is.
async fn run_for_node(
    config: &Config,
    node: &Node,
    cmds: &[(&'static str, PreServeCommand)],
    dry_run: bool,
) -> Result<()> {
    for (origin, cmd) in cmds {
        let on_failure = cmd.on_failure();
        let (label_cmd, ssh_cmd) = match cmd {
            PreServeCommand::Script { script, .. } => {
                let body = crate::builtin_scripts::resolve(script).ok_or_else(|| {
                    anyhow::anyhow!(
                        "Unknown built-in script alias '{}' on node {}. Known: {}",
                        script,
                        node.name,
                        crate::builtin_scripts::known_aliases()
                    )
                })?;
                let remote = format!("/tmp/sm121-builtin-{script}.sh");
                crate::ssh::upload_script(config, node, body, &remote, dry_run).await?;
                (
                    format!("script:{script}"),
                    format!("sudo bash {}", shq(&remote)),
                )
            }
            _ => {
                let c = cmd.command();
                (c.to_string(), c.to_string())
            }
        };
        if dry_run {
            info!(
                "{}",
                fluent_i18n::t!("after-stop-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!(
                        "after_stop command failed on {}: {}",
                        node.name,
                        e
                    ));
                }
                OnFailure::Warn => warn!(
                    "{}",
                    fluent_i18n::t!("after-stop-ignored", { "cmd" => &node.name, "error" => e.to_string() })
                ),
                OnFailure::Continue => {}
            }
        }
    }
    Ok(())
}

fn truncate(s: &str, max: usize) -> String {
    // Same char-safe truncate as before_serve (multibyte-safe).
    let head: String = s.chars().take(max).collect();
    if s.chars().count() <= max {
        s.to_string()
    } else {
        format!("{head}…")
    }
}

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

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

    fn test_config() -> Config {
        Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![test_cluster()],
            nccl: std::collections::HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: std::collections::HashMap::new(),
            exec_aliases: std::collections::HashMap::new(),
            home_assistant: None,
            mikrotik: None,
            drop_caches: true,
            mount_aliases: std::collections::HashMap::new(),
            recipe_sources: Vec::new(),
            docker_image_sources: Vec::new(),
            cluster_sources: Vec::new(),
            registry: None,
            global_mods_dir: None,
        }
    }

    fn test_recipe(after_stop: Vec<PreServeCommand>) -> Recipe {
        Recipe {
            build_profile: None,
            parent_model: None,
            archived: false,
            engine: "vllm".to_string(),
            name: "T".to_string(),
            docker_image: String::new(),
            model: "x".to_string(),
            port: 30000,
            required_nodes: None,
            pipeline_parallel_size: None,
            tensor_parallel_size: 1,
            args: crate::pre_serve::RecipeArgs::default(),
            env: std::collections::HashMap::new(),
            cluster: None,
            entrypoint: vec![],
            distributed_backend: Some(crate::config::DistributedBackend::Mp),
            before_serve: vec![],
            after_stop,
            solo: None,
            deploy: crate::config::DeployTarget::Docker,
            venv: None,
            venv_setup: VenvSetup::default(),
            venv_setup_script: None,
            mods: Default::default(),
            sync_image: true,
            docker_image_dir: None,
            shm_size: String::new(),
            extra_mounts: vec![],
            build: Default::default(),
            composed_model: None,
        }
    }

    #[tokio::test]
    async fn empty_after_stop_does_nothing() {
        let cfg = test_config();
        let c = cfg.clusters[0].clone();
        let r = test_recipe(vec![]);
        // No hooks → immediate Ok, no SSH.
        super::run(&cfg, &c, &r, true).await.unwrap();
    }

    #[tokio::test]
    async fn dry_run_emits_for_each_command() {
        let cfg = test_config();
        let c = cfg.clusters[0].clone();
        let r = test_recipe(vec![
            PreServeCommand::Bare("echo done".into()),
            PreServeCommand::Script {
                script: "restore_dgx_services".into(),
                on_failure: OnFailure::Warn,
            },
        ]);
        // Dry-run must not error on the known alias.
        super::run(&cfg, &c, &r, true).await.unwrap();
    }

    #[test]
    fn deserialize_after_stop_script_entry() {
        let toml_str = r#"
name = "Demo"
model = "x"
after_stop = [
    "echo cleaned up",
    { script = "restore_dgx_services" },
    { script = "restore_dgx_services", on_failure = "warn" },
]
"#;
        let r: Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(r.after_stop.len(), 3);
        assert!(matches!(r.after_stop[0], PreServeCommand::Bare(_)));
        assert!(matches!(r.after_stop[1], PreServeCommand::Script { .. }));
        // Known alias resolves.
        assert!(crate::builtin_scripts::resolve("restore_dgx_services").is_some());
    }

    #[test]
    fn deserialize_before_and_after_stop_coexist() {
        let toml_str = r#"
name = "Demo"
model = "x"
before_serve = [{ script = "stop_dgx_services" }]
after_stop   = [{ script = "restore_dgx_services" }]
"#;
        let r: Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(r.before_serve.len(), 1);
        assert_eq!(r.after_stop.len(), 1);
    }
}