|
|
@@ -0,0 +1,263 @@
|
|
|
+//! `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(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!(
|
|
|
+ "{}",
|
|
|
+ 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};
|
|
|
+
|
|
|
+ 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,
|
|
|
+ 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,
|
|
|
+ drop_caches: true,
|
|
|
+ mount_aliases: std::collections::HashMap::new(),
|
|
|
+ global_mods_dir: None,
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ fn test_recipe(after_stop: Vec<PreServeCommand>) -> Recipe {
|
|
|
+ Recipe {
|
|
|
+ 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::config::RecipeArgs::default(),
|
|
|
+ env: std::collections::HashMap::new(),
|
|
|
+ cluster: None,
|
|
|
+ entrypoint: vec![],
|
|
|
+ distributed_backend: Some(crate::config::DistributedBackend::Mp),
|
|
|
+ before_serve: vec![],
|
|
|
+ after_stop,
|
|
|
+ mods: Default::default(),
|
|
|
+ sync_image: true,
|
|
|
+ docker_image_dir: None,
|
|
|
+ shm_size: String::new(),
|
|
|
+ extra_mounts: vec![],
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ #[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);
|
|
|
+ }
|
|
|
+}
|
