diff --git a/src/cloudai/workloads/sleep/sleep.py b/src/cloudai/workloads/sleep/sleep.py index 216981314..69c46882d 100644 --- a/src/cloudai/workloads/sleep/sleep.py +++ b/src/cloudai/workloads/sleep/sleep.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"); @@ -15,9 +15,11 @@ # limitations under the License. -from cloudai.core import Installable +from cloudai.core import Installable, JobStatusResult, TestRun from cloudai.models.workload import CmdArgs, TestDefinition +EXIT_CODE_FILE_NAME = "exit_code.txt" + class SleepCmdArgs(CmdArgs): """Sleep test command arguments.""" @@ -34,3 +36,37 @@ class SleepTestDefinition(TestDefinition): @property def installables(self) -> list[Installable]: return [] + + def was_run_successful(self, tr: TestRun) -> JobStatusResult: + exit_code_path = tr.output_path / EXIT_CODE_FILE_NAME + if not exit_code_path.is_file(): + # Not every system's command-gen strategy captures an exit code yet + # (currently only standalone does); fall back to the previous + # behavior rather than fail runs that never had a chance to produce + # this file. + return JobStatusResult(is_successful=True) + + try: + exit_code_text = exit_code_path.read_text().strip() + except (OSError, UnicodeDecodeError) as e: + return JobStatusResult( + is_successful=False, + error_message=f"Could not read exit code file {exit_code_path}: {e}.", + ) + + 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: + stderr_path = tr.output_path / "stderr.txt" + return JobStatusResult( + is_successful=False, + error_message=f"sleep exited with code {exit_code}. Check {stderr_path} for details.", + ) + + return JobStatusResult(is_successful=True) diff --git a/src/cloudai/workloads/sleep/standalone_command_gen_strategy.py b/src/cloudai/workloads/sleep/standalone_command_gen_strategy.py index b4494b971..9f5b05762 100644 --- a/src/cloudai/workloads/sleep/standalone_command_gen_strategy.py +++ b/src/cloudai/workloads/sleep/standalone_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 import toml @@ -21,7 +22,7 @@ from cloudai.core import CommandGenStrategy from cloudai.models.scenario import TestRunDetails -from .sleep import SleepTestDefinition +from .sleep import EXIT_CODE_FILE_NAME, SleepTestDefinition class SleepStandaloneCommandGenStrategy(CommandGenStrategy): @@ -33,11 +34,21 @@ def _generate_sleep_command(self) -> str: def store_test_run(self) -> None: test_cmd = self._generate_sleep_command() + full_cmd = self.gen_exec_command(store=False) self.test_run.output_path.mkdir(parents=True, exist_ok=True) with (self.test_run.output_path / self.TEST_RUN_DUMP_FILE_NAME).open("w", encoding="utf-8") as f: - trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=test_cmd) + trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=full_cmd) toml.dump(trd.model_dump(), f) - def gen_exec_command(self) -> str: - self.store_test_run() - return self._generate_sleep_command() + def gen_exec_command(self, *, store: bool = True) -> str: + if store: + self.store_test_run() + stdout_path = self.test_run.output_path / "stdout.txt" + stderr_path = self.test_run.output_path / "stderr.txt" + exit_code_path = self.test_run.output_path / EXIT_CODE_FILE_NAME + + return ( + f"{self._generate_sleep_command()} " + f"> {shlex.quote(str(stdout_path))} 2> {shlex.quote(str(stderr_path))}; " + f"echo $? > {shlex.quote(str(exit_code_path))}" + ) diff --git a/tests/workloads/sleep/test_command_gen_strategy_standalone.py b/tests/workloads/sleep/test_command_gen_strategy_standalone.py index 95964f422..ad63cd89d 100644 --- a/tests/workloads/sleep/test_command_gen_strategy_standalone.py +++ b/tests/workloads/sleep/test_command_gen_strategy_standalone.py @@ -22,6 +22,7 @@ from cloudai.models.scenario import TestRunDetails from cloudai.systems.standalone import StandaloneSystem from cloudai.workloads.sleep import SleepCmdArgs, SleepStandaloneCommandGenStrategy, SleepTestDefinition +from cloudai.workloads.sleep.sleep import EXIT_CODE_FILE_NAME def test_gen_exec_command_writes_test_run_details(tmp_path: Path, standalone_system: StandaloneSystem) -> None: @@ -36,10 +37,13 @@ def test_gen_exec_command_writes_test_run_details(tmp_path: Path, standalone_sys strategy = SleepStandaloneCommandGenStrategy(standalone_system, tr) command = strategy.gen_exec_command() - assert command == "sleep 60" + assert command == ( + f"sleep 60 > {tr.output_path / 'stdout.txt'} 2> {tr.output_path / 'stderr.txt'}; " + f"echo $? > {tr.output_path / EXIT_CODE_FILE_NAME}" + ) dump_path = tr.output_path / CommandGenStrategy.TEST_RUN_DUMP_FILE_NAME assert dump_path.is_file() details = TestRunDetails.model_validate(toml.load(dump_path)) assert details.test_cmd == "sleep 60" - assert details.full_cmd == "sleep 60" + assert details.full_cmd == command diff --git a/tests/workloads/sleep/test_sleep.py b/tests/workloads/sleep/test_sleep.py new file mode 100644 index 000000000..0382d2c79 --- /dev/null +++ b/tests/workloads/sleep/test_sleep.py @@ -0,0 +1,96 @@ +# 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 +from unittest.mock import patch + +from cloudai.core import TestRun +from cloudai.workloads.sleep import SleepCmdArgs, SleepTestDefinition +from cloudai.workloads.sleep.sleep import EXIT_CODE_FILE_NAME + + +def _sleep_test_run(tmp_path: Path) -> TestRun: + tdef = SleepTestDefinition( + name="sleep_test", + description="Simple sleep test", + test_template_name="Sleep", + cmd_args=SleepCmdArgs(seconds=1), + ) + return TestRun(name="sleep-job", test=tdef, num_nodes=1, nodes=[], output_path=tmp_path / "output") + + +def test_was_run_successful_without_exit_code_file_assumes_success(tmp_path: Path) -> None: + """Systems that don't yet capture an exit code (Slurm, LSF, Kubernetes) keep the old behavior.""" + tr = _sleep_test_run(tmp_path) + + result = tr.test.was_run_successful(tr) + + assert result.is_successful is True + + +def test_was_run_successful_with_zero_exit_code(tmp_path: Path) -> None: + tr = _sleep_test_run(tmp_path) + tr.output_path.mkdir(parents=True) + (tr.output_path / EXIT_CODE_FILE_NAME).write_text("0\n") + + result = tr.test.was_run_successful(tr) + + assert result.is_successful is True + + +def test_was_run_successful_with_nonzero_exit_code(tmp_path: Path) -> None: + tr = _sleep_test_run(tmp_path) + tr.output_path.mkdir(parents=True) + (tr.output_path / EXIT_CODE_FILE_NAME).write_text("1\n") + + result = tr.test.was_run_successful(tr) + + assert result.is_successful is False + assert "exited with code 1" in result.error_message + + +def test_was_run_successful_with_unparseable_exit_code(tmp_path: Path) -> None: + tr = _sleep_test_run(tmp_path) + tr.output_path.mkdir(parents=True) + (tr.output_path / EXIT_CODE_FILE_NAME).write_text("not-a-number\n") + + result = tr.test.was_run_successful(tr) + + assert result.is_successful is False + assert "Could not parse exit code" in result.error_message + + +def test_was_run_successful_with_unreadable_exit_code_file(tmp_path: Path) -> None: + tr = _sleep_test_run(tmp_path) + tr.output_path.mkdir(parents=True) + (tr.output_path / EXIT_CODE_FILE_NAME).write_text("0\n") + + with patch("pathlib.Path.read_text", side_effect=PermissionError("permission denied")): + result = tr.test.was_run_successful(tr) + + assert result.is_successful is False + assert "Could not read exit code file" in result.error_message + + +def test_was_run_successful_with_undecodable_exit_code_file(tmp_path: Path) -> None: + tr = _sleep_test_run(tmp_path) + tr.output_path.mkdir(parents=True) + (tr.output_path / EXIT_CODE_FILE_NAME).write_bytes(b"\xff\xfe\x00") + + result = tr.test.was_run_successful(tr) + + assert result.is_successful is False + assert "Could not read exit code file" in result.error_message