From d6d97b210d44fb62891eeef243bf2081b4e67085 Mon Sep 17 00:00:00 2001 From: Marcin Malagowski Date: Wed, 8 Jul 2026 13:51:13 +0300 Subject: [PATCH 1/2] Report run status from command exit code in SlurmContainer Grade the run via was_run_successful() - now based on the per-test exit_code.txt file (non-zero = FAILED). Document the exit-code file and its parsing and failure rules. Signed-off-by: Marcin Malagowski --- doc/workloads/slurm_container.rst | 12 ++- .../slurm_command_gen_strategy.py | 14 ++- .../slurm_container/slurm_container.py | 39 +++++++- tests/ref_data/slurm_container.sbatch | 2 +- .../test_command_gen_strategy_slurm.py | 44 +++++++-- .../slurm_container/test_slurm_container.py | 89 +++++++++++++++++++ 6 files changed, 187 insertions(+), 13 deletions(-) create mode 100644 tests/workloads/slurm_container/test_slurm_container.py diff --git a/doc/workloads/slurm_container.rst b/doc/workloads/slurm_container.rst index 718329139..bf50ce8e8 100644 --- a/doc/workloads/slurm_container.rst +++ b/doc/workloads/slurm_container.rst @@ -15,7 +15,7 @@ Test TOML example: test_template_name = "SlurmContainer" [cmd_args] - image_path = "/path/to/container.sqsh" + docker_image_url = "/path/to/container.sqsh" cmd = "python train.py" Test Scenario example: @@ -47,9 +47,17 @@ Test-in-Scenario example: test_template_name = "SlurmContainer" [Tests.cmd_args] - image_path = "/path/to/container.sqsh" + docker_image_url = "/path/to/container.sqsh" cmd = "python train.py" +Run Status +---------- + +CloudAI determines success by reading the per-test ``exit_code.txt`` file from +the test output directory. Only integer ``0`` is successful. Any non-zero +integer, or a missing, malformed, unreadable, or undecodable file, marks the run +as failed. Slurm-style exit code ``0:0`` is not accepted. + API Documentation ----------------- diff --git a/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py b/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py index 2f20fbb53..b21f38449 100644 --- a/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.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"); @@ -18,7 +18,7 @@ from cloudai.systems.slurm import SlurmCommandGenStrategy -from .slurm_container import SlurmContainerTestDefinition +from .slurm_container import EXIT_CODE_FILE_NAME, SlurmContainerTestDefinition class SlurmContainerCommandGenStrategy(SlurmCommandGenStrategy): @@ -46,4 +46,12 @@ def generate_test_command(self) -> list[str]: if self.test_run.test.extra_cmd_args: command_parts.append(self.test_run.test.extra_args_str) - return command_parts + command = " ".join(command_parts) + wrapped_command = ( + f"( {command} ); " + r"rc=\$?; " + rf"printf '%s\n' \"\$rc\" > /cloudai_run_results/{EXIT_CODE_FILE_NAME}; " + r"exit \"\$rc\"" + ) + + return [wrapped_command] diff --git a/src/cloudai/workloads/slurm_container/slurm_container.py b/src/cloudai/workloads/slurm_container/slurm_container.py index b81c85c49..ad189f703 100644 --- a/src/cloudai/workloads/slurm_container/slurm_container.py +++ b/src/cloudai/workloads/slurm_container/slurm_container.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"); @@ -18,9 +18,11 @@ from pydantic import Field -from cloudai.core import DockerImage, File, Installable +from cloudai.core import DockerImage, File, Installable, JobStatusResult, TestRun from cloudai.models.workload import CmdArgs, TestDefinition +EXIT_CODE_FILE_NAME = "exit_code.txt" + class SlurmContainerCmdArgs(CmdArgs): """Command line arguments for a generic Slurm container test.""" @@ -53,3 +55,36 @@ def extra_args_str(self) -> str: for k, v in self.extra_cmd_args.items(): parts.append(f"{k} {v}" if v else k) return " ".join(parts) + + def was_run_successful(self, tr: TestRun) -> JobStatusResult: + """Grade the run from the container command exit code.""" + exit_code_path = tr.output_path / EXIT_CODE_FILE_NAME + if not exit_code_path.is_file(): + return JobStatusResult( + is_successful=False, + error_message=f"Exit code file {exit_code_path} not found.", + ) + + try: + exit_code_text = exit_code_path.read_text(encoding="utf-8").strip() + except (OSError, UnicodeDecodeError) as err: + return JobStatusResult( + is_successful=False, + error_message=f"Could not read exit code file {exit_code_path}: {err}.", + ) + + try: + exit_code = int(exit_code_text) + except ValueError: + return JobStatusResult( + is_successful=False, + error_message=f"Could not parse exit code from {exit_code_path}: {exit_code_text!r}.", + ) + + if exit_code != 0: + return JobStatusResult( + is_successful=False, + error_message=f"Container command exited with code {exit_code}.", + ) + + return JobStatusResult(is_successful=True) diff --git a/tests/ref_data/slurm_container.sbatch b/tests/ref_data/slurm_container.sbatch index 586cad22c..e85855dc5 100644 --- a/tests/ref_data/slurm_container.sbatch +++ b/tests/ref_data/slurm_container.sbatch @@ -14,4 +14,4 @@ srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --containe srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh -srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; pwd ; ls" +srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; ( pwd ; ls ); rc=\$?; printf '%s\n' \"\$rc\" > /cloudai_run_results/exit_code.txt; exit \"\$rc\"" diff --git a/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py b/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py index 2a6b1b3f9..2eaff6b77 100644 --- a/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py +++ b/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py @@ -18,14 +18,24 @@ import pytest -from cloudai.core import TestRun +from cloudai.core import TestRun, TestScenario from cloudai.models.workload import NsysConfiguration -from cloudai.systems.slurm import SlurmSystem +from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem from cloudai.workloads.slurm_container import ( SlurmContainerCmdArgs, SlurmContainerCommandGenStrategy, SlurmContainerTestDefinition, ) +from cloudai.workloads.slurm_container.slurm_container import EXIT_CODE_FILE_NAME + + +def _wrapped_command(command: str) -> str: + return ( + f"( {command} ); " + r"rc=\$?; " + rf"printf '%s\n' \"\$rc\" > /cloudai_run_results/{EXIT_CODE_FILE_NAME}; " + r"exit \"\$rc\"" + ) @pytest.fixture @@ -52,7 +62,9 @@ def test_default(slurm_system: SlurmSystem, test_run: TestRun) -> None: f"--no-container-mount-home" ) - assert cmd == f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; cmd"' + assert cmd == ( + f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; {_wrapped_command("cmd")}"' + ) def test_with_nsys(slurm_system: SlurmSystem, test_run: TestRun) -> None: @@ -70,7 +82,10 @@ def test_with_nsys(slurm_system: SlurmSystem, test_run: TestRun) -> None: f"--no-container-mount-home" ) - assert cmd == f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; nsys profile cmd"' + assert cmd == ( + f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; ' + f'{_wrapped_command("nsys profile cmd")}"' + ) def test_with_extra_srun_args(slurm_system: SlurmSystem, test_run: TestRun) -> None: @@ -91,4 +106,23 @@ def test_with_extra_srun_args(slurm_system: SlurmSystem, test_run: TestRun) -> N f"{' '.join(extra_args)}" ) - assert cmd == f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; cmd"' + assert cmd == ( + f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; {_wrapped_command("cmd")}"' + ) + + +def test_single_sbatch_writes_exit_code_to_per_test_output(slurm_system: SlurmSystem, test_run: TestRun) -> None: + test_run.output_path = slurm_system.output_path / "single-batch" + test_run.output_path.mkdir(parents=True) + scenario = TestScenario(name="tc", test_runs=[test_run]) + runner = SingleSbatchRunner( + mode="run", + system=slurm_system, + test_scenario=scenario, + output_path=slurm_system.output_path, + ) + + block = runner.get_single_tr_block(test_run) + + assert f"{test_run.output_path.absolute()}:/cloudai_run_results" in block + assert f"/cloudai_run_results/{EXIT_CODE_FILE_NAME}" in block diff --git a/tests/workloads/slurm_container/test_slurm_container.py b/tests/workloads/slurm_container/test_slurm_container.py new file mode 100644 index 000000000..33cdf7920 --- /dev/null +++ b/tests/workloads/slurm_container/test_slurm_container.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import patch + +import pytest + +from cloudai.core import TestRun +from cloudai.workloads.slurm_container import SlurmContainerCmdArgs, SlurmContainerTestDefinition +from cloudai.workloads.slurm_container.slurm_container import EXIT_CODE_FILE_NAME + + +class TestSlurmContainerSuccessCheck: + def setup_method(self) -> None: + self.tdef = SlurmContainerTestDefinition( + name="sc", + description="desc", + test_template_name="SlurmContainer", + cmd_args=SlurmContainerCmdArgs(docker_image_url="docker://url", cmd="bash /scripts/run.sh"), + ) + + def _write_exit_code(self, tr: TestRun, exit_code: str) -> None: + tr.output_path.mkdir(parents=True, exist_ok=True) + (tr.output_path / EXIT_CODE_FILE_NAME).write_text(exit_code, encoding="utf-8") + + def test_missing_exit_code_fails(self, base_tr: TestRun) -> None: + result = self.tdef.was_run_successful(base_tr) + + assert not result.is_successful + assert EXIT_CODE_FILE_NAME in result.error_message + assert "not found" in result.error_message + + @pytest.mark.parametrize( + ("exit_code", "is_successful"), + [ + ("0", True), + ("0\n", True), + ("1", False), + ("42", False), + ("137", False), + ], + ) + def test_exit_code_is_honored(self, base_tr: TestRun, exit_code: str, is_successful: bool) -> None: + self._write_exit_code(base_tr, exit_code) + + result = self.tdef.was_run_successful(base_tr) + + assert result.is_successful is is_successful + if not is_successful: + assert exit_code.strip() in result.error_message + + def test_malformed_exit_code_is_reported(self, base_tr: TestRun) -> None: + self._write_exit_code(base_tr, "not-a-number") + + result = self.tdef.was_run_successful(base_tr) + + assert not result.is_successful + assert "Could not parse exit code" in result.error_message + + def test_unreadable_exit_code_is_reported(self, base_tr: TestRun) -> None: + self._write_exit_code(base_tr, "0") + + with patch("pathlib.Path.read_text", side_effect=PermissionError("permission denied")): + result = self.tdef.was_run_successful(base_tr) + + assert not result.is_successful + assert "Could not read exit code file" in result.error_message + + def test_undecodable_exit_code_is_reported(self, base_tr: TestRun) -> None: + base_tr.output_path.mkdir(parents=True, exist_ok=True) + (base_tr.output_path / EXIT_CODE_FILE_NAME).write_bytes(b"\xff\xfe\x00") + + result = self.tdef.was_run_successful(base_tr) + + assert not result.is_successful + assert "Could not read exit code file" in result.error_message From 3cd366f58a9aea5efe9240ad235b90467c069dc7 Mon Sep 17 00:00:00 2001 From: Marcin Malagowski Date: Tue, 4 Aug 2026 18:27:40 +0300 Subject: [PATCH 2/2] Record aggregate srun status for SlurmContainer Move exit-code recording outside per-task container execution so one host-side writer records the aggregate srun result. Add multi-task coverage and update the sbatch reference. Signed-off-by: Marcin Malagowski --- .../slurm_command_gen_strategy.py | 21 +++++----- tests/ref_data/slurm_container.sbatch | 2 +- .../test_command_gen_strategy_slurm.py | 38 ++++++++++++++----- 3 files changed, 41 insertions(+), 20 deletions(-) diff --git a/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py b/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py index b21f38449..3eb0cd4b4 100644 --- a/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/slurm_container/slurm_command_gen_strategy.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import shlex from typing import cast from cloudai.systems.slurm import SlurmCommandGenStrategy @@ -40,18 +41,20 @@ def gen_srun_prefix(self, use_pretest_extras: bool = False, with_num_nodes: bool tdef: SlurmContainerTestDefinition = cast(SlurmContainerTestDefinition, self.test_run.test) return [*cmd, *tdef.extra_srun_args] + def _gen_srun_command(self) -> str: + srun_command = super()._gen_srun_command() + exit_code_path = shlex.quote(str((self.test_run.output_path / EXIT_CODE_FILE_NAME).absolute())) + return ( + f"{srun_command}; " + "rc=$?; " + f"""printf '%s\\n' "$rc" > {exit_code_path}; """ + """(exit "$rc")""" + ) + def generate_test_command(self) -> list[str]: tdef: SlurmContainerTestDefinition = cast(SlurmContainerTestDefinition, self.test_run.test) command_parts: list[str] = [*super().gen_nsys_command(), tdef.cmd_args.cmd] if self.test_run.test.extra_cmd_args: command_parts.append(self.test_run.test.extra_args_str) - command = " ".join(command_parts) - wrapped_command = ( - f"( {command} ); " - r"rc=\$?; " - rf"printf '%s\n' \"\$rc\" > /cloudai_run_results/{EXIT_CODE_FILE_NAME}; " - r"exit \"\$rc\"" - ) - - return [wrapped_command] + return command_parts diff --git a/tests/ref_data/slurm_container.sbatch b/tests/ref_data/slurm_container.sbatch index e85855dc5..702aa9389 100644 --- a/tests/ref_data/slurm_container.sbatch +++ b/tests/ref_data/slurm_container.sbatch @@ -14,4 +14,4 @@ srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --containe srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output --ntasks=1 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash /cloudai_install/slurm-metadata.sh -srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; ( pwd ; ls ); rc=\$?; printf '%s\n' \"\$rc\" > /cloudai_run_results/exit_code.txt; exit \"\$rc\"" +srun --export=ALL --mpi=pmix -N1 --container-image=https://docker/url --container-mounts=__OUTPUT_DIR__/output:/cloudai_run_results,__OUTPUT_DIR__/install:/cloudai_install,__OUTPUT_DIR__/output bash -c "source __OUTPUT_DIR__/output/env_vars.sh; pwd ; ls"; rc=$?; printf '%s\n' "$rc" > __OUTPUT_DIR__/output/exit_code.txt; (exit "$rc") diff --git a/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py b/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py index 2eaff6b77..ab86f48cc 100644 --- a/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py +++ b/tests/workloads/slurm_container/test_command_gen_strategy_slurm.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import shlex from typing import cast import pytest @@ -29,12 +30,12 @@ from cloudai.workloads.slurm_container.slurm_container import EXIT_CODE_FILE_NAME -def _wrapped_command(command: str) -> str: +def _status_capture(test_run: TestRun) -> str: + exit_code_path = shlex.quote(str((test_run.output_path / EXIT_CODE_FILE_NAME).absolute())) return ( - f"( {command} ); " - r"rc=\$?; " - rf"printf '%s\n' \"\$rc\" > /cloudai_run_results/{EXIT_CODE_FILE_NAME}; " - r"exit \"\$rc\"" + "; rc=$?; " + f"""printf '%s\\n' "$rc" > {exit_code_path}; """ + """(exit "$rc")""" ) @@ -63,7 +64,8 @@ def test_default(slurm_system: SlurmSystem, test_run: TestRun) -> None: ) assert cmd == ( - f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; {_wrapped_command("cmd")}"' + f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; cmd"' + f"{_status_capture(test_run)}" ) @@ -83,8 +85,8 @@ def test_with_nsys(slurm_system: SlurmSystem, test_run: TestRun) -> None: ) assert cmd == ( - f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; ' - f'{_wrapped_command("nsys profile cmd")}"' + f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; nsys profile cmd"' + f"{_status_capture(test_run)}" ) @@ -107,7 +109,8 @@ def test_with_extra_srun_args(slurm_system: SlurmSystem, test_run: TestRun) -> N ) assert cmd == ( - f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; {_wrapped_command("cmd")}"' + f'{srun_part} bash -c "source {(test_run.output_path / "env_vars.sh").absolute()}; cmd"' + f"{_status_capture(test_run)}" ) @@ -125,4 +128,19 @@ def test_single_sbatch_writes_exit_code_to_per_test_output(slurm_system: SlurmSy block = runner.get_single_tr_block(test_run) assert f"{test_run.output_path.absolute()}:/cloudai_run_results" in block - assert f"/cloudai_run_results/{EXIT_CODE_FILE_NAME}" in block + assert str((test_run.output_path / EXIT_CODE_FILE_NAME).absolute()) in block + assert f"/cloudai_run_results/{EXIT_CODE_FILE_NAME}" not in block + + +def test_multi_task_run_records_one_aggregate_srun_status(slurm_system: SlurmSystem, test_run: TestRun) -> None: + tdef = cast(SlurmContainerTestDefinition, test_run.test) + tdef.extra_srun_args = ["--ntasks=2"] + tdef.cmd_args.cmd = r"bash -c 'exit \$SLURM_PROCID'" + cgs = SlurmContainerCommandGenStrategy(slurm_system, test_run) + + cmd = cgs.gen_srun_command() + exit_code_path = str((test_run.output_path / EXIT_CODE_FILE_NAME).absolute()) + + assert "--ntasks=2" in cmd + assert r"""bash -c 'exit \$SLURM_PROCID'"; rc=$?;""" in cmd + assert cmd.count(exit_code_path) == 1