|
|
@@ -22,6 +22,7 @@ pub async fn run(
|
|
|
config: &Config,
|
|
|
cluster: &Cluster,
|
|
|
recipe: &Recipe,
|
|
|
+ recipe_name: &str,
|
|
|
skip: bool,
|
|
|
dry_run: bool,
|
|
|
) -> Result<()> {
|
|
|
@@ -41,6 +42,13 @@ pub async fn run(
|
|
|
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());
|
|
|
@@ -48,6 +56,7 @@ pub async fn run(
|
|
|
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 {
|
|
|
(
|
|
|
@@ -57,7 +66,15 @@ pub async fn run(
|
|
|
// 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, dry).await,
|
|
|
+ run_for_node(
|
|
|
+ &config,
|
|
|
+ &node,
|
|
|
+ local_rank as u32,
|
|
|
+ &cmds,
|
|
|
+ dir.as_deref(),
|
|
|
+ dry,
|
|
|
+ )
|
|
|
+ .await,
|
|
|
)
|
|
|
}));
|
|
|
}
|
|
|
@@ -111,6 +128,7 @@ async fn run_for_node(
|
|
|
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 {
|
|
|
@@ -120,20 +138,29 @@ async fn run_for_node(
|
|
|
// command (`Bare`/`Configured`) runs as-is.
|
|
|
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-{}.sh", script);
|
|
|
- crate::ssh::upload_script(config, node, body, &remote, dry_run).await?;
|
|
|
+ // 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:{script}"),
|
|
|
+ format!("script:{tag}"),
|
|
|
format!("sudo NODE_RANK={local_rank} bash {remote}"),
|
|
|
)
|
|
|
}
|
|
|
@@ -192,268 +219,34 @@ fn truncate(s: &str, max: usize) -> String {
|
|
|
}
|
|
|
}
|
|
|
|
|
|
-#[cfg(test)]
|
|
|
-mod tests {
|
|
|
- use super::*;
|
|
|
- use crate::config::{Node, PreServeCommand};
|
|
|
- use crate::pre_serve::VenvSetup;
|
|
|
- use std::collections::HashMap;
|
|
|
-
|
|
|
- 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,
|
|
|
- },
|
|
|
- Node {
|
|
|
- name: "n1".to_string(),
|
|
|
- ssh_ip: "10.0.0.2".to_string(),
|
|
|
- ib_ip: Some("1.2.3.5".to_string()),
|
|
|
- rank: 1,
|
|
|
- host_key: None,
|
|
|
- entity: None,
|
|
|
- },
|
|
|
- ],
|
|
|
- extra_entities: vec![],
|
|
|
- }
|
|
|
- }
|
|
|
-
|
|
|
- fn test_config(before_serve: HashMap<String, crate::config::BeforeServeTable>) -> Config {
|
|
|
- Config {
|
|
|
- ssh_user: "u".to_string(),
|
|
|
- ssh_key_path: "/k".into(),
|
|
|
- clusters: vec![test_cluster()],
|
|
|
- nccl: HashMap::new(),
|
|
|
- hf_cache_dir: None,
|
|
|
- init_port: 25000,
|
|
|
- before_serve,
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
- home_assistant: None,
|
|
|
- mikrotik: None,
|
|
|
- global_mods_dir: None,
|
|
|
- drop_caches: true,
|
|
|
- mount_aliases: HashMap::new(),
|
|
|
- recipe_sources: Vec::new(),
|
|
|
- recipe_repos: Vec::new(),
|
|
|
- docker_image_sources: Vec::new(),
|
|
|
- cluster_sources: Vec::new(),
|
|
|
- registry: None,
|
|
|
- }
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn truncate_keeps_short_string_unchanged() {
|
|
|
- assert_eq!(truncate("short", 120), "short");
|
|
|
- assert_eq!(truncate("", 120), "");
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn truncate_shortens_to_char_boundary() {
|
|
|
- // 10 chars, cap at 4 → first 4 chars + ellipsis
|
|
|
- let out = truncate("abcdefghij", 4);
|
|
|
- assert_eq!(out, "abcd…");
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn truncate_does_not_panic_on_multibyte() {
|
|
|
- // Each 'é' is 2 bytes; a byte slice `&s[..3]` would land mid-char and
|
|
|
- // panic. The char-based impl must not. (Regression for C-2.)
|
|
|
- let s = "éééééé"; // 6 chars, 12 bytes
|
|
|
- assert_eq!(truncate(s, 3), "ééé…");
|
|
|
- // cap exactly at a multibyte boundary must also be safe
|
|
|
- assert_eq!(truncate(s, 4), "éééé…");
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn truncate_cjk_chars() {
|
|
|
- // 3-byte CJK: byte-slice at any non-multiple-of-3 would panic.
|
|
|
- let s = "你好世界你好";
|
|
|
- assert_eq!(truncate(s, 2), "你好…");
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn combined_runs_global_then_cluster_then_recipe() {
|
|
|
- let mut bs = HashMap::new();
|
|
|
- bs.insert(
|
|
|
- "c1".to_string(),
|
|
|
- crate::config::BeforeServeTable {
|
|
|
- commands: vec![PreServeCommand::Bare("cluster_cmd".into())],
|
|
|
- },
|
|
|
- );
|
|
|
- bs.insert(
|
|
|
- "*".to_string(),
|
|
|
- crate::config::BeforeServeTable {
|
|
|
- commands: vec![PreServeCommand::Bare("global_cmd".into())],
|
|
|
- },
|
|
|
- );
|
|
|
- let cfg = test_config(bs);
|
|
|
- let recipe = Recipe {
|
|
|
- before_serve: vec![PreServeCommand::Bare("recipe_cmd".into())],
|
|
|
- ..test_recipe()
|
|
|
- };
|
|
|
-
|
|
|
- let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
|
|
|
- let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
|
|
|
- let origins: Vec<&str> = cmds.iter().map(|(o, _)| *o).collect();
|
|
|
- assert_eq!(names, vec!["global_cmd", "cluster_cmd", "recipe_cmd"]);
|
|
|
- assert_eq!(origins, vec!["global", "cluster", "recipe"]);
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn combined_global_applies_to_unlisted_cluster() {
|
|
|
- let mut bs = HashMap::new();
|
|
|
- bs.insert(
|
|
|
- "*".to_string(),
|
|
|
- crate::config::BeforeServeTable {
|
|
|
- commands: vec![PreServeCommand::Bare("global_cmd".into())],
|
|
|
- },
|
|
|
- );
|
|
|
- // Cluster "c1" has no specific entry — only the global hook runs.
|
|
|
- let cfg = test_config(bs);
|
|
|
- let recipe = Recipe {
|
|
|
- before_serve: vec![],
|
|
|
- after_stop: vec![],
|
|
|
- solo: None,
|
|
|
- ..test_recipe()
|
|
|
- };
|
|
|
-
|
|
|
- let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
|
|
|
- let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
|
|
|
- assert_eq!(names, vec!["global_cmd"]);
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn combined_orders_cluster_then_recipe() {
|
|
|
- let mut bs = HashMap::new();
|
|
|
- bs.insert(
|
|
|
- "c1".to_string(),
|
|
|
- crate::config::BeforeServeTable {
|
|
|
- commands: vec![PreServeCommand::Bare("cluster_cmd_1".into())],
|
|
|
- },
|
|
|
+/// 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()
|
|
|
);
|
|
|
- let cfg = test_config(bs);
|
|
|
- let recipe = Recipe {
|
|
|
- before_serve: vec![
|
|
|
- PreServeCommand::Bare("recipe_cmd_1".into()),
|
|
|
- PreServeCommand::Bare("recipe_cmd_2".into()),
|
|
|
- ],
|
|
|
- ..test_recipe()
|
|
|
- };
|
|
|
-
|
|
|
- let cmds = combined_commands(&cfg, &test_cluster(), &recipe);
|
|
|
- let names: Vec<&str> = cmds.iter().map(|(_, c)| c.command()).collect();
|
|
|
- let origins: Vec<&str> = cmds.iter().map(|(o, _)| *o).collect();
|
|
|
- assert_eq!(names, vec!["cluster_cmd_1", "recipe_cmd_1", "recipe_cmd_2"]);
|
|
|
- assert_eq!(origins, vec!["cluster", "recipe", "recipe"]);
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn combined_empty_when_unconfigured() {
|
|
|
- let cfg = test_config(HashMap::new());
|
|
|
- let recipe = Recipe {
|
|
|
- before_serve: vec![],
|
|
|
- after_stop: vec![],
|
|
|
- solo: None,
|
|
|
- ..test_recipe()
|
|
|
- };
|
|
|
- assert!(combined_commands(&cfg, &test_cluster(), &recipe).is_empty());
|
|
|
- }
|
|
|
-
|
|
|
- fn test_recipe() -> Recipe {
|
|
|
- Recipe {
|
|
|
- build_profile: None,
|
|
|
- parent_model: None,
|
|
|
- archived: false,
|
|
|
- engine: "vllm".to_string(),
|
|
|
- name: "Test".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(),
|
|
|
- cold_args: crate::pre_serve::RecipeArgs::default(),
|
|
|
- env: HashMap::new(),
|
|
|
- cluster: None,
|
|
|
- entrypoint: vec![],
|
|
|
- distributed_backend: Some(crate::config::DistributedBackend::Mp),
|
|
|
- before_serve: vec![],
|
|
|
- after_stop: vec![],
|
|
|
- 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(),
|
|
|
- devices: vec![],
|
|
|
- memory_limit: None,
|
|
|
- cpus: None,
|
|
|
- extra_mounts: vec![],
|
|
|
- extra_mounts_manifest: None,
|
|
|
- build: Default::default(),
|
|
|
- composed_model: None,
|
|
|
- staging: vec![],
|
|
|
- dgpp: None,
|
|
|
- }
|
|
|
- }
|
|
|
-
|
|
|
- // ── Deserialization tests (the ones that live in this module) ──
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn parses_cluster_level_before_serve() {
|
|
|
- let toml_str = r#"
|
|
|
-ssh_user = "u"
|
|
|
-ssh_key_path = "/k"
|
|
|
-
|
|
|
-[[clusters]]
|
|
|
-name = "c1"
|
|
|
-master_ib_ip = "1.1.1.1"
|
|
|
-nodes = [
|
|
|
- { name = "n0", ssh_ip = "2.2.2.2", ib_ip = "1.1.1.1", rank = 0 },
|
|
|
-]
|
|
|
-
|
|
|
-[before_serve."c1"]
|
|
|
-commands = [
|
|
|
- "sudo nvidia-smi -lgc 300,2300",
|
|
|
- { command = "sudo nvidia-smi -pm 1", on_failure = "warn" },
|
|
|
-]
|
|
|
-"#;
|
|
|
- let cfg: Config = toml::from_str(toml_str).unwrap();
|
|
|
- assert!(cfg.before_serve.contains_key("c1"));
|
|
|
- let list = &cfg.before_serve["c1"].commands;
|
|
|
- assert_eq!(list.len(), 2);
|
|
|
- assert_eq!(list[0].command(), "sudo nvidia-smi -lgc 300,2300");
|
|
|
- assert_eq!(list[0].on_failure(), OnFailure::Abort); // default
|
|
|
- assert_eq!(list[1].command(), "sudo nvidia-smi -pm 1");
|
|
|
- assert_eq!(list[1].on_failure(), OnFailure::Warn);
|
|
|
- }
|
|
|
-
|
|
|
- #[test]
|
|
|
- fn parses_recipe_level_before_serve() {
|
|
|
- let toml_str = r#"
|
|
|
-name = "Demo"
|
|
|
-model = "x"
|
|
|
-before_serve = [
|
|
|
- "echo a",
|
|
|
- { command = "echo b", on_failure = "continue" },
|
|
|
-]
|
|
|
-"#;
|
|
|
- let r: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
- assert_eq!(r.before_serve.len(), 2);
|
|
|
- assert_eq!(r.before_serve[0].command(), "echo a");
|
|
|
- assert_eq!(r.before_serve[0].on_failure(), OnFailure::Abort);
|
|
|
- assert_eq!(r.before_serve[1].command(), "echo b");
|
|
|
- assert_eq!(r.before_serve[1].on_failure(), OnFailure::Continue);
|
|
|
- }
|
|
|
+ };
|
|
|
+ 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;
|
