From 99eed934435302251487a9912faf3a487e43c8cc Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 15 Aug 2026 21:22:24 +0200 Subject: [PATCH 1/5] pin-nodes + training uses correct nodes allocated --- src/cloudai/_core/test_scenario.py | 1 + src/cloudai/models/scenario.py | 7 ++++ .../report_generator/training/parser.py | 25 +++++++++++- .../systems/slurm/single_sbatch_runner.py | 2 + src/cloudai/systems/slurm/slurm_job.py | 6 +-- src/cloudai/systems/slurm/slurm_metadata.py | 6 ++- src/cloudai/systems/slurm/slurm_runner.py | 22 +++++++++- src/cloudai/systems/slurm/slurm_system.py | 5 ++- src/cloudai/test_scenario_parser.py | 1 + .../training/test_training_parser.py | 40 +++++++++++++++++++ tests/systems/slurm/test_allocation.py | 3 +- tests/test_get_job_id.py | 27 ++++++++++++- tests/test_test_scenario.py | 13 +++++- 13 files changed, 144 insertions(+), 14 deletions(-) diff --git a/src/cloudai/_core/test_scenario.py b/src/cloudai/_core/test_scenario.py index 32fa25c62..b8bdded3a 100644 --- a/src/cloudai/_core/test_scenario.py +++ b/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)) diff --git a/src/cloudai/models/scenario.py b/src/cloudai/models/scenario.py index b910c0a6a..45869f33c 100644 --- a/src/cloudai/models/scenario.py +++ b/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 = { diff --git a/src/cloudai/report_generator/training/parser.py b/src/cloudai/report_generator/training/parser.py index 0447b3527..24756945f 100644 --- a/src/cloudai/report_generator/training/parser.py +++ b/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 metadata.num_nodes or 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") diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 5e7110fce..c8d51b608 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/src/cloudai/systems/slurm/single_sbatch_runner.py @@ -279,4 +279,6 @@ 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, + num_nodes=len(job.nodes) if job.nodes else None, ) diff --git a/src/cloudai/systems/slurm/slurm_job.py b/src/cloudai/systems/slurm/slurm_job.py index 834c74373..c5633b765 100644 --- a/src/cloudai/systems/slurm/slurm_job.py +++ b/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) diff --git a/src/cloudai/systems/slurm/slurm_metadata.py b/src/cloudai/systems/slurm/slurm_metadata.py index c06eab691..3446f3b83 100644 --- a/src/cloudai/systems/slurm/slurm_metadata.py +++ b/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,8 @@ class SlurmJobMetadata(_SlurmStepMetadataBase): is_single_sbatch: bool = False job_root: Path job_steps: list[SlurmStepMetadata] + nodes: list[str] = Field(default_factory=list) + num_nodes: int | None = None @field_serializer("job_root") def _path_serializer(self, v: Path) -> str: diff --git a/src/cloudai/systems/slurm/slurm_runner.py b/src/cloudai/systems/slurm/slurm_runner.py index dae0cdb29..b51dc97a4 100644 --- a/src/cloudai/systems/slurm/slurm_runner.py +++ b/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,17 @@ 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: + raise RuntimeError(f"Cannot pin test case '{tr.name}': its first job has no recorded node allocation") + 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 +135,8 @@ 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, + num_nodes=len(job.nodes) if job.nodes else None, ) def store_job_metadata(self, job: SlurmJob): diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 145c32056..edef2989b 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/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 diff --git a/src/cloudai/test_scenario_parser.py b/src/cloudai/test_scenario_parser.py index b469c3709..ae5c0c73e 100644 --- a/src/cloudai/test_scenario_parser.py +++ b/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), diff --git a/tests/report_generator/training/test_training_parser.py b/tests/report_generator/training/test_training_parser.py index 3fb17b53a..e641bfe54 100644 --- a/tests/report_generator/training/test_training_parser.py +++ b/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,44 @@ 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"], + num_nodes=2, + ) + 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}} diff --git a/tests/systems/slurm/test_allocation.py b/tests/systems/slurm/test_allocation.py index c8f06a624..0e672b382 100644 --- a/tests/systems/slurm/test_allocation.py +++ b/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) diff --git a/tests/test_get_job_id.py b/tests/test_get_job_id.py index ecdf6ced3..28bf83044 100644 --- a/tests/test_get_job_id.py +++ b/tests/test_get_job_id.py @@ -95,14 +95,39 @@ 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(tr.name, tr.test, 2, [], 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) + + @pytest.mark.parametrize( "stdout, stderr, expected_job_id", [ diff --git a/tests/test_test_scenario.py b/tests/test_test_scenario.py index 64eb2c844..7c02c1267 100644 --- a/tests/test_test_scenario.py +++ b/tests/test_test_scenario.py @@ -267,7 +267,9 @@ 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( @@ -275,6 +277,7 @@ def test_create_test_run_with_hooks(tdef: TestDefinition, test_scenario_parser: ) 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 ): From e3f2996bccdd509376f7c5f298a2cfa7a1883f59 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 15 Aug 2026 21:30:54 +0200 Subject: [PATCH 2/5] make one test more meaningful --- tests/test_get_job_id.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_get_job_id.py b/tests/test_get_job_id.py index 28bf83044..bf85ea121 100644 --- a/tests/test_get_job_id.py +++ b/tests/test_get_job_id.py @@ -116,7 +116,7 @@ def test_slurm_runner_records_and_reuses_nodes_per_case(slurm_runner: SlurmRunne assert slurm_runner.pinned_nodes == {tr.name: ["node01", "node02"]} - next_tr = TestRun(tr.name, tr.test, 2, [], pin_nodes=True) + 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)) From 515c608c1cc6a9172d107a62fe40ef1eb052a941 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Sat, 15 Aug 2026 21:34:15 +0200 Subject: [PATCH 3/5] recover from not being able to pin nodes --- src/cloudai/systems/slurm/slurm_runner.py | 7 ++++--- tests/test_get_job_id.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/cloudai/systems/slurm/slurm_runner.py b/src/cloudai/systems/slurm/slurm_runner.py index b51dc97a4..374115f12 100644 --- a/src/cloudai/systems/slurm/slurm_runner.py +++ b/src/cloudai/systems/slurm/slurm_runner.py @@ -96,9 +96,10 @@ def on_job_completion(self, job: BaseJob) -> None: 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: - raise RuntimeError(f"Cannot pin test case '{tr.name}': its first job has no recorded node allocation") - self.pinned_nodes[tr.name] = slurm_job.nodes.copy() - logging.info("Pinned test case '%s' to nodes: %s", tr.name, ",".join(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: diff --git a/tests/test_get_job_id.py b/tests/test_get_job_id.py index bf85ea121..e57ee984e 100644 --- a/tests/test_get_job_id.py +++ b/tests/test_get_job_id.py @@ -128,6 +128,25 @@ def test_slurm_runner_records_and_reuses_nodes_per_case(slurm_runner: SlurmRunne 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() + + @pytest.mark.parametrize( "stdout, stderr, expected_job_id", [ From 7212cc2c2ebf1ed46e7cb10c9b0028d870b445d0 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 17 Aug 2026 20:55:59 +0200 Subject: [PATCH 4/5] derive training node count from allocation --- src/cloudai/report_generator/training/parser.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cloudai/report_generator/training/parser.py b/src/cloudai/report_generator/training/parser.py index 24756945f..4f85f0b4a 100644 --- a/src/cloudai/report_generator/training/parser.py +++ b/src/cloudai/report_generator/training/parser.py @@ -215,7 +215,7 @@ def _get_used_nodes(tr: TestRun) -> tuple[int, list[str]]: if not metadata.nodes: return tr.nnodes, list(tr.nodes) - return metadata.num_nodes or len(metadata.nodes), list(metadata.nodes) + return len(metadata.nodes), list(metadata.nodes) @staticmethod def _get_clique_size(env_vars: dict[str, Any]) -> Optional[int]: From 7e75ba8774b9bde961534074ad6a47d12ea2c366 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 18 Aug 2026 14:17:09 +0200 Subject: [PATCH 5/5] remove redundant num_nodes field --- src/cloudai/systems/slurm/single_sbatch_runner.py | 1 - src/cloudai/systems/slurm/slurm_metadata.py | 1 - src/cloudai/systems/slurm/slurm_runner.py | 1 - tests/report_generator/training/test_training_parser.py | 1 - 4 files changed, 4 deletions(-) diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index c8d51b608..3a49041f4 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/src/cloudai/systems/slurm/single_sbatch_runner.py @@ -280,5 +280,4 @@ def _get_job_metadata( is_single_sbatch=True, job_root=self.scenario_root.absolute(), nodes=job.nodes, - num_nodes=len(job.nodes) if job.nodes else None, ) diff --git a/src/cloudai/systems/slurm/slurm_metadata.py b/src/cloudai/systems/slurm/slurm_metadata.py index 3446f3b83..a3d2e9ed8 100644 --- a/src/cloudai/systems/slurm/slurm_metadata.py +++ b/src/cloudai/systems/slurm/slurm_metadata.py @@ -81,7 +81,6 @@ class SlurmJobMetadata(_SlurmStepMetadataBase): job_root: Path job_steps: list[SlurmStepMetadata] nodes: list[str] = Field(default_factory=list) - num_nodes: int | None = None @field_serializer("job_root") def _path_serializer(self, v: Path) -> str: diff --git a/src/cloudai/systems/slurm/slurm_runner.py b/src/cloudai/systems/slurm/slurm_runner.py index 374115f12..1980211ee 100644 --- a/src/cloudai/systems/slurm/slurm_runner.py +++ b/src/cloudai/systems/slurm/slurm_runner.py @@ -137,7 +137,6 @@ def _get_job_metadata( test_cmd=" ".join(cmd_gen.generate_test_command()), job_root=job.test_run.output_path.absolute(), nodes=job.nodes, - num_nodes=len(job.nodes) if job.nodes else None, ) def store_job_metadata(self, job: SlurmJob): diff --git a/tests/report_generator/training/test_training_parser.py b/tests/report_generator/training/test_training_parser.py index e641bfe54..c20f82847 100644 --- a/tests/report_generator/training/test_training_parser.py +++ b/tests/report_generator/training/test_training_parser.py @@ -291,7 +291,6 @@ def test_build_config_uses_nodes_from_slurm_job_metadata(tmp_path: Path): job_root=tmp_path, job_steps=[], nodes=["node01", "node02"], - num_nodes=2, ) with (tmp_path / "slurm-job.toml").open("w") as f: toml.dump(metadata.model_dump(mode="json"), f)