[Slurm, Reporting] pin-nodes + proper nodes storage

src/cloudai/_core/test_scenario.py

@@ -101,6 +101,7 @@ class TestRun:
     reports: Set[Type[ReportGenerationStrategy]] = field(default_factory=set)
     extra_srun_args: str | None = None
     num_nodes_explicit: bool = False
+    pin_nodes: bool = False
 
     def __hash__(self) -> int:
         return hash(self.name + self.test.name + str(self.iterations) + str(self.current_iteration))

src/cloudai/models/scenario.py

@@ -74,6 +74,7 @@ class TestRunModel(BaseModel):
     test_name: Optional[str] = None
     num_nodes: int | list[int] | None = None
     nodes: list[str] = Field(default_factory=list)
+    pin_nodes: bool = False
     exclude_nodes: list[str] = Field(
         default_factory=list,
         description=(
@@ -112,6 +113,12 @@ def parse_sol(cls, value: Any) -> float | cloudai.metrics.MetricSOLConfig | None
             return cloudai.metrics.parse_sol_spec(value)
         return value
 
+    @model_validator(mode="after")
+    def validate_pin_nodes(self) -> Self:
+        if self.pin_nodes and isinstance(self.num_nodes, list):
+            raise ValueError("pin_nodes cannot be enabled when num_nodes is swept")
+        return self
+
     def tdef_model_dump(self, by_alias: bool) -> dict:
         """Return a dictionary with non-None values that correspond to the test definition fields."""
         data = {

src/cloudai/report_generator/training/parser.py

@@ -29,9 +29,11 @@
 from pathlib import Path
 from typing import Any, ClassVar, Optional
 
+import toml
 import yaml
 
 from cloudai.core import System, TestRun, TestScenario
+from cloudai.systems.slurm import SlurmJobMetadata
 
 from .mappings import (
     MEGATRON_BRIDGE_MODEL_CONFIG,
@@ -160,6 +162,7 @@ def _build_config(self, tr: TestRun, system: System, test_scenario: TestScenario
         """Map the framework + test config into TrainingConfig, then fill the CloudAI-computed fields."""
         env_vars = {**getattr(system, "global_env_vars", {}), **tr.test.extra_env_vars}
         config_paths = test_scenario.config_paths
+        num_nodes, nodes = self._get_used_nodes(tr)
         config = TrainingConfig(
             test_id=tr.name,
             test_name=tr.test.name,
@@ -171,8 +174,8 @@ def _build_config(self, tr: TestRun, system: System, test_scenario: TestScenario
             test_scenario_path=str(config_paths.test_scenario_path) if config_paths is not None else "",
             cloudai_execution_node=socket.gethostname(),
             env_vars=env_vars,
-            num_nodes=tr.nnodes,
-            nodes=list(tr.nodes),
+            num_nodes=num_nodes,
+            nodes=nodes,
             **self._resolve_model_config(tr),
             **self._resolve_test_config(tr),
         )
@@ -196,6 +199,24 @@ def _build_config(self, tr: TestRun, system: System, test_scenario: TestScenario
         config.model_name = self.get_model_name(tr)
         return config
 
+    @staticmethod
+    def _get_used_nodes(tr: TestRun) -> tuple[int, list[str]]:
+        slurm_job_path = tr.output_path / "slurm-job.toml"
+        if not slurm_job_path.is_file():
+            return tr.nnodes, list(tr.nodes)
+
+        try:
+            with slurm_job_path.open() as f:
+                metadata = SlurmJobMetadata.model_validate(toml.load(f))
+        except Exception as exc:
+            logging.warning("Could not load Slurm allocation from '%s': %s", slurm_job_path, exc)
+            return tr.nnodes, list(tr.nodes)
+
+        if not metadata.nodes:
+            return tr.nnodes, list(tr.nodes)
+
+        return len(metadata.nodes), list(metadata.nodes)
+
     @staticmethod
     def _get_clique_size(env_vars: dict[str, Any]) -> Optional[int]:
         clique_size = env_vars.get("CLIQUE_SIZE")

src/cloudai/systems/slurm/single_sbatch_runner.py

@@ -279,4 +279,5 @@ def _get_job_metadata(
             test_cmd="n/a for single sbatch run",
             is_single_sbatch=True,
             job_root=self.scenario_root.absolute(),
+            nodes=job.nodes,
         )

src/cloudai/systems/slurm/slurm_job.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from cloudai.core import BaseJob
 
@@ -23,4 +23,4 @@
 class SlurmJob(BaseJob):
     """A job class for execution on a Slurm system."""
 
-    pass
+    nodes: list[str] = field(default_factory=list, init=False)

src/cloudai/systems/slurm/slurm_metadata.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,7 +18,7 @@
 
 from pathlib import Path
 
-from pydantic import BaseModel, ConfigDict, field_serializer
+from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
 
 class _SlurmStepMetadataBase(BaseModel):
@@ -80,6 +80,7 @@ class SlurmJobMetadata(_SlurmStepMetadataBase):
     is_single_sbatch: bool = False
     job_root: Path
     job_steps: list[SlurmStepMetadata]
+    nodes: list[str] = Field(default_factory=list)
 
     @field_serializer("job_root")
     def _path_serializer(self, v: Path) -> str:

src/cloudai/systems/slurm/slurm_runner.py

@@ -42,6 +42,13 @@ def __init__(self, mode: str, system: System, test_scenario: TestScenario, outpu
         super().__init__(mode, system, test_scenario, output_path)
         self.system = cast(SlurmSystem, system)
         self.cmd_shell = CommandShell()
+        self.pinned_nodes: dict[str, list[str]] = {}
+
+    def submit_test(self, tr: TestRun) -> None:
+        if tr.pin_nodes and tr.name in self.pinned_nodes:
+            tr.nodes = self.pinned_nodes[tr.name].copy()
+            logging.info("Forcing test case '%s' to use pinned nodes: %s", tr.name, ",".join(tr.nodes))
+        super().submit_test(tr)
 
     def get_job_id(self, stdout: str, stderr: str) -> int | None:
         match = re.search(r"Submitted batch job (\d+)", stdout)
@@ -82,8 +89,18 @@ def completed_test_runs(self, job: BaseJob) -> list[TestRun]:
 
     def on_job_completion(self, job: BaseJob) -> None:
         logging.debug(f"Job completion callback for job {job.id}")
-        self.system.complete_job(cast(SlurmJob, job))
-        self.store_job_metadata(cast(SlurmJob, job))
+        slurm_job = cast(SlurmJob, job)
+        slurm_job.nodes = self.system.complete_job(slurm_job)
+        self.store_job_metadata(slurm_job)
+
+        tr = slurm_job.test_run
+        if self.mode == "run" and tr.pin_nodes and tr.name not in self.pinned_nodes:
+            if not slurm_job.nodes:
+                logging.error("Cannot pin test case '%s': the job has no recorded node allocation", tr.name)
+            else:
+                self.pinned_nodes[tr.name] = slurm_job.nodes.copy()
+                logging.info("Pinned test case '%s' to nodes: %s", tr.name, ",".join(slurm_job.nodes))
+
         for tr in self.completed_test_runs(job):
             try:
                 self.get_cmd_gen_strategy(self.system, tr).cleanup_job_artifacts()
@@ -119,6 +136,7 @@ def _get_job_metadata(
             srun_cmd=cmd_gen.gen_srun_command(),
             test_cmd=" ".join(cmd_gen.generate_test_command()),
             job_root=job.test_run.output_path.absolute(),
+            nodes=job.nodes,
         )
 
     def store_job_metadata(self, job: SlurmJob):

src/cloudai/systems/slurm/slurm_system.py

@@ -805,9 +805,10 @@ def get_nodes_by_spec(
     def system_installables(self) -> list[Installable]:
         return [File(Path(__file__).parent.absolute() / "slurm-metadata.sh")]
 
-    def complete_job(self, job: SlurmJob) -> None:
+    def complete_job(self, job: SlurmJob) -> list[str]:
         out, _ = self.fetch_command_output(f"sacct -j {job.id} -p --noheader -X --format=NodeList")
         spec = out.splitlines()[0] if out.splitlines() else out
-        nodelist = set(parse_node_list(spec.strip().replace("|", "")))
+        nodelist = sorted(set(parse_node_list(spec.strip().replace("|", ""))))
         to_unlock = [node for node in self.group_allocated if node.name in nodelist]
         self.group_allocated.difference_update(to_unlock)
+        return nodelist

src/cloudai/test_scenario_parser.py

@@ -208,6 +208,7 @@ def _create_test_run(
             num_nodes=test_info.num_nodes or 1,
             iterations=test_info.iterations,
             nodes=test_info.nodes,
+            pin_nodes=test_info.pin_nodes,
             time_limit=total_time_limit,
             sol=legacy_sol,
             metric_sol=cloudai.metrics.merge_sol_configs(self.system.sol, scenario_sol, metric_sol),

tests/report_generator/training/test_training_parser.py

@@ -21,6 +21,7 @@
 from typing import Any
 
 import pytest
+import toml
 
 from cloudai.core import ConfigPaths
 from cloudai.models.scenario import ReportConfig
@@ -30,6 +31,7 @@
 from cloudai.report_generator.training.models import SCHEMA_VERSION, Scalar, TrainingResults, TrainingStep
 from cloudai.report_generator.training.parser import MegatronBridgeParser, MegatronParser, NeMoRunParser
 from cloudai.report_generator.training.reporter import TrainingReporter
+from cloudai.systems.slurm import SlurmJobMetadata
 
 
 def _scalars(rows: list[tuple]) -> list[Scalar]:
@@ -275,6 +277,43 @@ def test_build_config_resolves_paths_and_computes_fields():
     assert config.data_parallel_size == 8  # 32 / (tp4 * pp1 * cp1)
 
 
+def test_build_config_uses_nodes_from_slurm_job_metadata(tmp_path: Path):
+    metadata = SlurmJobMetadata(
+        job_id=123,
+        name="training",
+        state="COMPLETED",
+        start_time="",
+        end_time="",
+        elapsed_time_sec=1,
+        exit_code="0:0",
+        srun_cmd="srun training",
+        test_cmd="training",
+        job_root=tmp_path,
+        job_steps=[],
+        nodes=["node01", "node02"],
+    )
+    with (tmp_path / "slurm-job.toml").open("w") as f:
+        toml.dump(metadata.model_dump(mode="json"), f)
+
+    parser = NeMoRunParser()
+    parser.get_model_config = lambda tr: {
+        "parallelism": {
+            "tensor_model_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "context_parallel_size": 1,
+        }
+    }
+    config = parser._build_config(
+        _tr(output_path=tmp_path, nnodes=1, nodes=[], recipe_name="gpt3"),
+        _system(gpus_per_node=4),
+        _scenario(),
+    )
+
+    assert config.nodes == ["node01", "node02"]
+    assert config.num_nodes == 2
+    assert config.world_size == 8
+
+
 def test_build_config_leaves_world_size_none_without_gpus_per_node():
     # No gpus_per_node/ntasks_per_node on the system: world_size/data_parallel_size stay None, rest still resolves.
     raw = {"parallelism": {"tensor_model_parallel_size": 4, "pipeline_model_parallel_size": 1}}

tests/systems/slurm/test_allocation.py

@@ -95,9 +95,10 @@ def test_completion_clears_group_allocation_state(self, slurm_system: SlurmSyste
             "cloudai.systems.slurm.slurm_system.SlurmSystem.fetch_command_output",
             return_value=(f"{','.join(nodes_list)}|", ""),
         ):
-            system.complete_job(SlurmJob(id=1, test_run=Mock()))
+            completed_nodes = system.complete_job(SlurmJob(id=1, test_run=Mock()))
 
         assert len(system.group_allocated) == 0
+        assert completed_nodes == nodes_list
 
     def test_group_allocation_is_preserved_on_updated(self, slurm_system: SlurmSystem, monkeypatch: pytest.MonkeyPatch):
         system, all_nodes, _ = self.prepare(slurm_system, [], monkeypatch)

tests/test_get_job_id.py

@@ -95,11 +95,55 @@ def test_slurm_runner_on_job_completion_calls_cleanup(slurm_runner: SlurmRunner)
     cleanup = Mock()
     slurm_runner.get_cmd_gen_strategy = Mock(return_value=Mock(cleanup_job_artifacts=cleanup))
 
-    with patch.object(SlurmSystem, "complete_job") as complete_job:
+    with patch.object(SlurmSystem, "complete_job", return_value=["node01", "node02"]) as complete_job:
         slurm_runner.on_job_completion(job)
 
     complete_job.assert_called_once_with(job)
     slurm_runner.store_job_metadata.assert_called_once_with(job)
+    assert job.nodes == ["node01", "node02"]
+    cleanup.assert_called_once()
+
+
+def test_slurm_runner_records_and_reuses_nodes_per_case(slurm_runner: SlurmRunner, caplog: pytest.LogCaptureFixture):
+    tr = slurm_runner.test_scenario.test_runs[0]
+    tr.pin_nodes = True
+    first_job = SlurmJob(tr, id=1)
+    slurm_runner.store_job_metadata = Mock()
+    slurm_runner.get_cmd_gen_strategy = Mock(return_value=Mock(cleanup_job_artifacts=Mock()))
+
+    with patch.object(SlurmSystem, "complete_job", return_value=["node01", "node02"]):
+        slurm_runner.on_job_completion(first_job)
+
+    assert slurm_runner.pinned_nodes == {tr.name: ["node01", "node02"]}
+
+    next_tr = TestRun(name=tr.name, test=tr.test, num_nodes=2, nodes=["bla1", "bla2"], pin_nodes=True)
+    slurm_runner.on_job_submit = Mock()
+    slurm_runner._submit_test = Mock(return_value=SlurmJob(next_tr, id=2))
+
+    with caplog.at_level("INFO"):
+        slurm_runner.submit_test(next_tr)
+
+    assert next_tr.nodes == ["node01", "node02"]
+    assert "Forcing test case 'tr-name' to use pinned nodes: node01,node02" in caplog.text
+    slurm_runner.on_job_submit.assert_called_once_with(next_tr)
+
+
+def test_slurm_runner_continues_when_pinned_case_has_no_recorded_nodes(
+    slurm_runner: SlurmRunner, caplog: pytest.LogCaptureFixture
+) -> None:
+    tr = slurm_runner.test_scenario.test_runs[0]
+    tr.pin_nodes = True
+    job = SlurmJob(tr, id=1)
+    slurm_runner.store_job_metadata = Mock()
+    cleanup = Mock()
+    slurm_runner.get_cmd_gen_strategy = Mock(return_value=Mock(cleanup_job_artifacts=cleanup))
+
+    with patch.object(SlurmSystem, "complete_job", return_value=[]), caplog.at_level("ERROR"):
+        slurm_runner.on_job_completion(job)
+
+    assert tr.name not in slurm_runner.pinned_nodes
+    assert "Cannot pin test case 'tr-name': the job has no recorded node allocation" in caplog.text
+    slurm_runner.store_job_metadata.assert_called_once_with(job)
     cleanup.assert_called_once()
 
 

tests/test_test_scenario.py

@@ -267,14 +267,17 @@ def test_create_test_run_with_hooks(tdef: TestDefinition, test_scenario_parser:
         test_runs=[TestRun(name="post1", test=tdef, num_nodes=1, nodes=[], time_limit="00:20:00", iterations=1)],
     )
 
-    test_info = TestRunModel(id="main1", test_name="test1", time_limit="01:00:00", weight=10, iterations=1, num_nodes=1)
+    test_info = TestRunModel(
+        id="main1", test_name="test1", time_limit="01:00:00", weight=10, iterations=1, num_nodes=1, pin_nodes=True
+    )
     test_scenario_parser.test_mapping = {"test1": tdef}
 
     test_run = test_scenario_parser._create_test_run(
         test_info=test_info, normalized_weight=1.0, pre_test=pre_test, post_test=post_test
     )
 
     assert test_run.time_limit == "01:50:00"  # Main + pre + post hooks
+    assert test_run.pin_nodes is True
 
 
 def test_total_time_limit_with_empty_hooks():
@@ -492,6 +495,14 @@ def test_num_nodes_can_be_list(self, test_scenario_parser: TestScenarioParser, s
         )
         assert model.tests[0].num_nodes == [1, 2]
 
+    def test_pin_nodes_rejects_num_nodes_sweep(self) -> None:
+        with pytest.raises(ValueError, match="pin_nodes cannot be enabled when num_nodes is swept"):
+            TestRunModel(id="1", test_name="nccl", num_nodes=[1, 2], pin_nodes=True)
+
+    def test_pin_nodes_accepts_fixed_num_nodes(self) -> None:
+        model = TestRunModel(id="1", test_name="nccl", num_nodes=2, pin_nodes=True)
+        assert model.pin_nodes is True
+
     def test_agent_metrics_preserved_from_test_definition(
         self, test_scenario_parser: TestScenarioParser, slurm_system: SlurmSystem
     ):