diff --git a/src/cloudai/workloads/slurm_container/slurm_container.py b/src/cloudai/workloads/slurm_container/slurm_container.py index b81c85c49..fdda0788d 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"); @@ -16,10 +16,12 @@ from typing import Optional -from pydantic import Field +import toml +from pydantic import Field, ValidationError -from cloudai.core import DockerImage, File, Installable +from cloudai.core import DockerImage, File, Installable, JobStatusResult, TestRun from cloudai.models.workload import CmdArgs, TestDefinition +from cloudai.systems.slurm import SlurmJobMetadata class SlurmContainerCmdArgs(CmdArgs): @@ -53,3 +55,38 @@ 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 real container exit code recorded by Slurm.""" + slurm_job_path = tr.output_path / "slurm-job.toml" + if not slurm_job_path.is_file(): + return JobStatusResult( + is_successful=False, + error_message=( + f"slurm-job.toml file not found in the specified output directory {tr.output_path}. " + "This file is required to determine the container exit status." + ), + ) + + try: + with slurm_job_path.open("r", encoding="utf-8") as file: + metadata = SlurmJobMetadata.model_validate(toml.load(file)) + except (OSError, toml.TomlDecodeError, ValidationError) as err: + return JobStatusResult( + is_successful=False, + error_message=( + f"Failed to read Slurm job metadata from {slurm_job_path} " + f"(malformed or partially written slurm-job.toml?): {err}" + ), + ) + + if metadata.exit_code == "0" or metadata.exit_code.startswith("0:"): + return JobStatusResult(is_successful=True) + + return JobStatusResult( + is_successful=False, + error_message=( + f"Container command exited with a non-zero exit code for {tr.output_path}: " + f"state={metadata.state}, exit_code={metadata.exit_code}." + ), + ) 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..000174f11 --- /dev/null +++ b/tests/workloads/slurm_container/test_slurm_container.py @@ -0,0 +1,90 @@ +# 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 pathlib import Path + +import pytest +import toml + +from cloudai.core import TestRun +from cloudai.systems.slurm import SlurmJobMetadata +from cloudai.workloads.slurm_container import SlurmContainerCmdArgs, SlurmContainerTestDefinition + + +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_slurm_metadata(self, output_path: Path, *, exit_code: str, state: str = "COMPLETED") -> None: + output_path.mkdir(parents=True, exist_ok=True) + with (output_path / "slurm-job.toml").open("w", encoding="utf-8") as file: + toml.dump( + SlurmJobMetadata( + job_id=123, + name="sc", + state=state, + exit_code=exit_code, + start_time="2026-07-08T01:36:18", + end_time="2026-07-08T01:44:17", + elapsed_time_sec=479, + srun_cmd="srun test", + test_cmd="bash /scripts/run.sh", + job_root=output_path, + job_steps=[], + ).model_dump(), + file, + ) + + def test_missing_slurm_metadata_fails(self, base_tr: TestRun) -> None: + result = self.tdef.was_run_successful(base_tr) + + assert not result.is_successful + assert "slurm-job.toml file not found" in result.error_message + + @pytest.mark.parametrize( + ("exit_code", "is_successful"), + [ + ("0:0", True), + ("0", True), + ("0:15", True), + ("1:0", False), + ("15:0", False), + ("42:0", False), + ("137:0", False), + ], + ) + def test_exit_code_is_honored(self, base_tr: TestRun, exit_code: str, is_successful: bool) -> None: + self._write_slurm_metadata(base_tr.output_path, exit_code=exit_code) + + result = self.tdef.was_run_successful(base_tr) + + assert result.is_successful is is_successful + if not is_successful: + assert exit_code in result.error_message + + def test_malformed_metadata_is_reported(self, base_tr: TestRun) -> None: + base_tr.output_path.mkdir(parents=True, exist_ok=True) + (base_tr.output_path / "slurm-job.toml").write_text("not = valid = toml =", encoding="utf-8") + + result = self.tdef.was_run_successful(base_tr) + + assert not result.is_successful + assert "slurm-job.toml" in result.error_message