From bf00dad8b1107344b903cabe17b82f701619725a Mon Sep 17 00:00:00 2001 From: Mihir Mahajan Date: Tue, 18 Aug 2026 17:26:34 +0200 Subject: [PATCH 1/3] feat(resources): typed `nodes` on base resources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `[[resources.components]]` has a typed `nodes`, base `[resources]` did not — so in a het recipe the trainer component could declare its node count directly while the fleet, declared in `[resources]`, had to reach for `sbatch_extra = ["--nodes=N"]`. The `sbatch_extra` docstring listed `--nodes` among the things the typed schema doesn't cover, which is the asymmetry rather than a justification for it. A multi-node fleet is the normal shape for RL runs here, not an edge case. Unset emits no directive, so single-node recipes render byte-identical scripts to before. Verified against the scheduler: a fleet of 2 nodes plus a 1-node GPU trainer allocated as one het job, `srun --het-group=0` reaching both fleet nodes (run_01a0157a86e872819d79e8b05888e496). Two unit tests cover the emission and that base and component counts stay independent. --- src/config.rs | 9 ++++++++- src/runner.rs | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/config.rs b/src/config.rs index ffb1e7d..d983e65 100644 --- a/src/config.rs +++ b/src/config.rs @@ -461,6 +461,12 @@ pub struct Resources { pub cpus: u32, pub mem: String, pub time: String, + /// `--nodes`. Defaults to 1 when unset, matching + /// [`HetComponent::nodes`] — a het recipe's fleet component is + /// declared here, and needed a `sbatch_extra` escape hatch before + /// this existed while its trainer component had a typed field. + #[serde(default)] + pub nodes: Option, pub partition: Option, pub qos: Option, pub account: Option, @@ -469,7 +475,7 @@ 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 + /// (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 @@ -518,6 +524,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..9329ce5 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,36 @@ 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 rather than defaulting to an explicit 1, so + // single-node recipes emit the same script as before. + 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!( From 8e8742243d24e443b0e0b9be858236ba3b058096 Mon Sep 17 00:00:00 2001 From: Mihir Mahajan Date: Tue, 18 Aug 2026 18:30:27 +0200 Subject: [PATCH 2/3] style(resources): satisfy rustfmt, trim the `nodes` doc to house style MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI's first step is `cargo fmt --all -- --check`, and it failed on the two one-line `assert!`s in `base_and_component_nodes_are_independent`: their argument lists are 62 chars, just past rustfmt's 60-char `fn_call_width`, so rustfmt wants them vertical. Formatted — every other assert in `het_tests` already has that shape. fmt runs before clippy, so nothing after it had executed in CI and both clippy steps were unverified; `cargo clippy --no-deps -- -D warnings` and its `--features ui` variant both pass clean, so no `#[allow]` was needed anywhere. `Resources::nodes` goes from four doc lines to one, matching its counterpart `HetComponent::nodes` (one line) rather than the escape-hatch fields around it. The `sbatch_extra` asymmetry it recounted is history and lives in the PR description; the doc keeps what a reader of the struct needs — the default, and that unset emits no directive at all, which is what makes existing single-node recipes render byte-identical scripts. Same trim on the two-line narration in `base_resources_nodes_is_typed`, where the module's own in-test comments are one-liners. Also rewraps the `sbatch_extra` paragraph, left ragged when `--nodes` came out of its list. No behaviour change; 41 tests pass under `--all-features`. Co-Authored-By: Claude Opus 5 --- src/config.rs | 9 +++------ src/runner.rs | 13 +++++++++---- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/config.rs b/src/config.rs index d983e65..0adf9fe 100644 --- a/src/config.rs +++ b/src/config.rs @@ -461,10 +461,7 @@ pub struct Resources { pub cpus: u32, pub mem: String, pub time: String, - /// `--nodes`. Defaults to 1 when unset, matching - /// [`HetComponent::nodes`] — a het recipe's fleet component is - /// declared here, and needed a `sbatch_extra` escape hatch before - /// this existed while its trainer component had a typed field. + /// `--nodes`. Defaults to 1 when unset, which emits no directive at all. #[serde(default)] pub nodes: Option, pub partition: Option, @@ -475,8 +472,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`, `--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. diff --git a/src/runner.rs b/src/runner.rs index 9329ce5..af6993e 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -2125,8 +2125,7 @@ time = "12:00:00" script.contains("#SBATCH --nodes=2"), "base nodes missing:\n{script}" ); - // Unset stays unset rather than defaulting to an explicit 1, so - // single-node recipes emit the same script as before. + // Unset stays unset, so existing recipes render unchanged. let plain = render(BASE); assert!( !plain.contains("#SBATCH --nodes"), @@ -2139,8 +2138,14 @@ time = "12:00:00" 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!( + 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, From 138e35522716a6f78683a88e68f970369eaaf52c Mon Sep 17 00:00:00 2001 From: Mihir Mahajan Date: Tue, 18 Aug 2026 18:34:01 +0200 Subject: [PATCH 3/3] style(resources): drop redundant serde(default) on nodes `Resources` carries a container-level `#[serde(default)]`, and no sibling Option field repeats it. --- src/config.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/config.rs b/src/config.rs index 0adf9fe..d4f8ef7 100644 --- a/src/config.rs +++ b/src/config.rs @@ -462,7 +462,6 @@ pub struct Resources { pub mem: String, pub time: String, /// `--nodes`. Defaults to 1 when unset, which emits no directive at all. - #[serde(default)] pub nodes: Option, pub partition: Option, pub qos: Option,