Report run status from Slurm exit code in SlurmContainer

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}."
+            ),
+        )

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