diff --git a/src/config.rs b/src/config.rs index ffb1e7d..d4f8ef7 100644 --- a/src/config.rs +++ b/src/config.rs @@ -461,6 +461,8 @@ pub struct Resources { pub cpus: u32, pub mem: String, pub time: String, + /// `--nodes`. Defaults to 1 when unset, which emits no directive at all. + pub nodes: Option, pub partition: Option, pub qos: Option, pub account: Option, @@ -469,8 +471,8 @@ pub struct Resources { /// waiting for SLURM's drain detection. pub exclude_nodes: Option, /// Escape hatch for sbatch directives the typed schema doesn't cover - /// (e.g. `--array`, `--nodes`, `--mail-type`, `--gpu-bind`). Each entry - /// is rendered verbatim as a separate `#SBATCH` line, *after* the + /// (e.g. `--array`, `--mail-type`, `--gpu-bind`). Each entry is + /// rendered verbatim as a separate `#SBATCH` line, *after* the /// typed directives. Use this only for things labctl can't model — /// don't override `--cpus-per-task` or `--time` here, those have /// dedicated fields and overriding them confuses the dispatcher. @@ -518,6 +520,7 @@ impl Default for Resources { cpus: 1, mem: "4GB".to_string(), time: "00:10:00".to_string(), + nodes: None, partition: None, qos: None, account: None, diff --git a/src/runner.rs b/src/runner.rs index ee46011..af6993e 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -1583,6 +1583,9 @@ fn render_script( if let Some(qos) = recipe.resources.qos.as_ref().or(cluster.slurm.qos.as_ref()) { script.push_str(&format!("#SBATCH --qos={qos}\n")); } + if let Some(nodes) = recipe.resources.nodes { + script.push_str(&format!("#SBATCH --nodes={nodes}\n")); + } if let Some(exclude) = recipe.resources.exclude_nodes.as_ref() { script.push_str(&format!("#SBATCH --exclude={exclude}\n")); } @@ -2115,6 +2118,41 @@ time = "12:00:00" ); } + #[test] + fn base_resources_nodes_is_typed() { + let script = render(&format!("{BASE}nodes = 2\n")); + assert!( + script.contains("#SBATCH --nodes=2"), + "base nodes missing:\n{script}" + ); + // Unset stays unset, so existing recipes render unchanged. + let plain = render(BASE); + assert!( + !plain.contains("#SBATCH --nodes"), + "unset nodes must not emit a directive:\n{plain}" + ); + } + + #[test] + fn base_and_component_nodes_are_independent() { + let script = render(&format!( + "{BASE}nodes = 2\n\n[[resources.components]]\ngpus = 4\ncpus = 16\nmem = \"128GB\"\nnodes = 1\n" + )); + assert!( + script.contains("#SBATCH --nodes=2"), + "fleet nodes:\n{script}" + ); + assert!( + script.contains("#SBATCH --nodes=1"), + "trainer nodes:\n{script}" + ); + assert_eq!( + script.matches("#SBATCH --nodes=").count(), + 2, + "one per component:\n{script}" + ); + } + #[test] fn components_emit_one_separator_and_own_directives() { let script = render(&format!(