Report run status from Slurm exit code in SlurmContainer

doc/workloads/slurm_container.rst

@@ -15,7 +15,8 @@ Test TOML example:
    test_template_name = "SlurmContainer"
 
    [cmd_args]
-   image_path = "/path/to/container.sqsh"
+   check_exit_code = true
+   docker_image_url = "/path/to/container.sqsh"
    cmd = "python train.py"
 
 Test Scenario example:
@@ -47,9 +48,20 @@ Test-in-Scenario example:
    test_template_name = "SlurmContainer"
 
      [Tests.cmd_args]
-     image_path = "/path/to/container.sqsh"
+     check_exit_code = true
+     docker_image_url = "/path/to/container.sqsh"
      cmd = "python train.py"
 
+Run Status
+----------
+
+Exit-code checking is disabled by default for backwards compatibility. Set
+``check_exit_code = true`` in ``cmd_args`` to record the aggregate ``srun`` exit
+code in the per-test ``exit_code.txt`` file and use it to determine success.
+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
 -----------------
 

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");
@@ -14,11 +14,12 @@
 # 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
 
-from .slurm_container import SlurmContainerTestDefinition
+from .slurm_container import EXIT_CODE_FILE_NAME, SlurmContainerTestDefinition
 
 
 class SlurmContainerCommandGenStrategy(SlurmCommandGenStrategy):
@@ -40,6 +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()
+        tdef = cast(SlurmContainerTestDefinition, self.test_run.test)
+        if not tdef.cmd_args.check_exit_code:
+            return 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]

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,15 +18,21 @@
 
 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."""
 
     docker_image_url: str
     cmd: str
+    check_exit_code: bool = Field(
+        default=False,
+        description="Whether to grade the run from the container command exit code.",
+    )
 
 
 class SlurmContainerTestDefinition(TestDefinition):
@@ -53,3 +59,39 @@ 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."""
+        if not self.cmd_args.check_exit_code:
+            return JobStatusResult(is_successful=True)
+
+        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)

tests/workloads/slurm_container/test_command_gen_strategy_slurm.py

@@ -14,18 +14,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import shlex
 from typing import cast
 
 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 _status_capture(test_run: TestRun) -> str:
+    exit_code_path = shlex.quote(str((test_run.output_path / EXIT_CODE_FILE_NAME).absolute()))
+    return (
+        "; rc=$?; "
+        f"""printf '%s\\n' "$rc" > {exit_code_path}; """
+        """(exit "$rc")"""
+    )
 
 
 @pytest.fixture
@@ -34,7 +45,7 @@ def test_run() -> TestRun:
         name="sc",
         description="desc",
         test_template_name="tt",
-        cmd_args=SlurmContainerCmdArgs(docker_image_url="docker://url", cmd="cmd"),
+        cmd_args=SlurmContainerCmdArgs(docker_image_url="docker://url", cmd="cmd", check_exit_code=True),
     )
     tr = TestRun(name="name", test=tdef, num_nodes=1, nodes=[])
     return tr
@@ -52,7 +63,20 @@ 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()}; cmd"'
+        f"{_status_capture(test_run)}"
+    )
+
+
+def test_exit_code_check_disabled_does_not_capture_status(slurm_system: SlurmSystem, test_run: TestRun) -> None:
+    tdef = cast(SlurmContainerTestDefinition, test_run.test)
+    tdef.cmd_args.check_exit_code = False
+    cgs = SlurmContainerCommandGenStrategy(slurm_system, test_run)
+
+    cmd = cgs.gen_srun_command()
+
+    assert EXIT_CODE_FILE_NAME not in cmd
 
 
 def test_with_nsys(slurm_system: SlurmSystem, test_run: TestRun) -> None:
@@ -70,7 +94,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()}; nsys profile cmd"'
+        f"{_status_capture(test_run)}"
+    )
 
 
 def test_with_extra_srun_args(slurm_system: SlurmSystem, test_run: TestRun) -> None:
@@ -91,4 +118,39 @@ 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()}; cmd"'
+        f"{_status_capture(test_run)}"
+    )
+
+
+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 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

tests/workloads/slurm_container/test_slurm_container.py

@@ -0,0 +1,107 @@
+# 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",
+                check_exit_code=True,
+            ),
+        )
+
+    def test_exit_code_check_disabled_by_default(self, base_tr: TestRun) -> None:
+        tdef = SlurmContainerTestDefinition(
+            name="sc",
+            description="desc",
+            test_template_name="SlurmContainer",
+            cmd_args=SlurmContainerCmdArgs(docker_image_url="docker://url", cmd="bash /scripts/run.sh"),
+        )
+        self._write_exit_code(base_tr, "42")
+
+        result = tdef.was_run_successful(base_tr)
+
+        assert getattr(tdef.cmd_args, "check_exit_code", None) is False
+        assert result.is_successful
+
+    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