diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 4aa0972de..70aa34200 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/src/cloudai/systems/slurm/single_sbatch_runner.py @@ -21,6 +21,8 @@ from pathlib import Path from typing import Generator, Optional, cast +import cloudai.models.output +import cloudai.output from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParams from cloudai.core import BaseJob, JobStatusResult, Registry, System, TestRun, TestScenario @@ -205,10 +207,10 @@ def run(self): finally: self.jobs.remove(job) - self.handle_dse() - self.on_job_completion(job) - self.update_run_output(job) + # A non-None result triggers final per-run inspection; each workload determines its own success. + self.update_run_output(job, JobStatusResult(is_successful=is_completed)) + self.handle_dse() def handle_dse(self): registry = Registry() @@ -221,31 +223,68 @@ def handle_dse(self): agent_config = agent_class.get_config_class()(**agent_config_data) gym = CloudAIGymEnv(tr, self, rewards=agent_config.rewards) - for idx, combination in enumerate(tr.all_combinations, start=1): - sampled_env_params = gym.params.sample(idx) if gym.params is not None else {} - next_tr = tr.apply_params_set(combination, env_params=sampled_env_params) - next_tr.step = idx - next_tr.output_path = self.get_job_output_path(next_tr) - - if not next_tr.test.constraint_check(next_tr, self.system): - continue - - gym.test_run = next_tr - observation = gym.get_observation() - reward = gym.compute_reward(observation) - gym.trajectory.append( - step=idx, - action=combination, - reward=reward, - observation=observation, - env_params=sampled_env_params, - ) + try: + for idx, combination in enumerate(tr.all_combinations, start=1): + sampled_env_params = gym.params.sample(idx) if gym.params is not None else {} + next_tr = tr.apply_params_set(combination, env_params=sampled_env_params) + next_tr.step = idx + next_tr.output_path = self.get_job_output_path(next_tr) + + if not next_tr.test.constraint_check(next_tr, self.system): + continue + + gym.test_run = next_tr + observation = gym.get_observation() + reward = gym.compute_reward(observation) + gym.trajectory.append( + step=idx, + action=combination, + reward=reward, + observation=observation, + env_params=sampled_env_params, + ) + finally: + gym.update_output() def completed_test_runs(self, job: BaseJob) -> list[TestRun]: return list(self.all_trs) - def get_run_output(self, job: BaseJob, tr: TestRun, result: JobStatusResult | None = None) -> None: - return None + def get_run_output( + self, job: BaseJob, tr: TestRun, result: JobStatusResult | None = None + ) -> cloudai.models.output.Run | None: + run_job = copy.copy(cast(SlurmJob, job)) + allocation_metadata = run_job.metadata + run_job.metadata = None + output_arg = f"--output={tr.output_path.absolute()}/stdout.txt" + steps = ( + [step for step in allocation_metadata.job_steps if output_arg in step.submit_line.split()] + if allocation_metadata is not None + else [] + ) + if result is not None: + if not steps and not (tr.output_path / "stdout.txt").exists(): + result = None + else: + try: + result = tr.test.was_run_successful(tr) + except Exception as exc: + logging.warning("Cannot determine output status for %s: %s", tr.output_path, exc) + return None + run = super().get_run_output(run_job, tr, result) + if run is None: + return None + if run.status == "failed" and any(step.state.startswith("CANCELLED") for step in steps): + run.status = "cancelled" + starts = [self._output_timestamp(step.start_time) for step in steps] + finishes = [self._output_timestamp(step.end_time) for step in steps] + if starts and all(start is not None for start in starts): + run.start = min(start for start in starts if start is not None) + if finishes and all(finish is not None for finish in finishes): + run.finish = max(finish for finish in finishes if finish is not None) + run.duration = cloudai.output.elapsed_seconds(run.start, run.finish) + if run.duration is None and len(steps) == 1: + run.duration = steps[0].elapsed_time_sec + return run def _submit_test(self, tr: TestRun) -> SlurmJob: with open(self.scenario_root / "cloudai_sbatch_script.sh", "w") as f: diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py index 4afe91b96..6dc71c8ee 100644 --- a/tests/test_single_sbatch_runner.py +++ b/tests/test_single_sbatch_runner.py @@ -15,6 +15,7 @@ # limitations under the License. import copy +import datetime import re from pathlib import Path from typing import Generator, Optional, cast @@ -24,9 +25,13 @@ import pytest import toml +import cloudai.metrics +import cloudai.models.output +from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParams, EnvParamSpec -from cloudai.core import Registry, System, TestRun, TestScenario +from cloudai.core import JobStatusResult, Registry, System, TestRun, TestScenario from cloudai.systems.slurm import SingleSbatchRunner, SlurmJob, SlurmJobMetadata, SlurmSystem +from cloudai.systems.slurm.slurm_metadata import SlurmStepMetadata from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition from cloudai.workloads.nccl_test.slurm_command_gen_strategy import NcclTestSlurmCommandGenStrategy from cloudai.workloads.sleep import SleepCmdArgs, SleepTestDefinition @@ -582,22 +587,112 @@ def submit_after_shutdown(_: TestRun) -> SlurmJob: assert runner.jobs == [] -def test_run_removes_completed_job_from_tracking(sleep_tr: TestRun, slurm_system: SlurmSystem) -> None: - tc = TestScenario(name="tc", test_runs=[sleep_tr]) +@pytest.mark.parametrize("status", ["completed", "failed", "cancelled", "unknown"]) +def test_run_removes_completed_job_from_tracking(sleep_tr: TestRun, slurm_system: SlurmSystem, status: str) -> None: + second_tr = copy.deepcopy(sleep_tr) + second_tr.name = "second" + tc = TestScenario(name="tc", test_runs=[sleep_tr, second_tr]) runner = SingleSbatchRunner(mode="run", system=slurm_system, test_scenario=tc, output_path=slurm_system.output_path) + runs = list(runner.all_trs) job = SlurmJob(sleep_tr, id=123) runner._submit_test = Mock(return_value=job) runner.handle_dse = Mock() - runner.on_job_completion = Mock() - + runner.on_job_completion = Mock(side_effect=runner.store_job_metadata) + allocation = SlurmStepMetadata( + job_id=123, + step_id="", + name="allocation", + state="CANCELLED", + exit_code="1:0", + start_time="2026-01-02T03:04:00Z", + end_time="2026-01-02T03:04:30Z", + elapsed_time_sec=30, + submit_line="sbatch run.sh", + ) + steps = [allocation] + for tr, step, start, end in [(sleep_tr, "3", 1, 4), (second_tr, "4", 5, 9), (second_tr, "5", 7, 12)]: + if tr is second_tr and status == "unknown": + continue + steps.append( + allocation.model_copy( + update={ + "step_id": step, + "submit_line": f"srun --output={tr.output_path.absolute()}/stdout.txt bash run.sh", + "state": "CANCELLED" if tr is second_tr and status == "cancelled" else "FAILED", + "start_time": f"2026-01-02T03:04:{start:02}Z", + "end_time": f"2026-01-02T03:04:{end:02}Z", + "elapsed_time_sec": end - start, + } + ) + ) with ( patch.object(SlurmSystem, "is_job_completed", return_value=True), + patch.object(SlurmSystem, "get_job_status", return_value=steps), patch.object(SlurmSystem, "kill") as kill, + patch.object( + SleepTestDefinition, + "was_run_successful", + side_effect=lambda tr: JobStatusResult( + is_successful=tr.name == sleep_tr.name or status == "completed", + ), + ), + patch.object( + SleepTestDefinition, + "metric_observations", + side_effect=lambda system, tr: [ + cloudai.metrics.MetricObservation(cloudai.metrics.BANDWIDTH, 5 if tr.name == sleep_tr.name else 10, {}), + ], + ), ): runner.run() + assert [run.duration for test in runner.experiment_output.experiment.tests for run in test.runs] == [ + 3, + None if status == "unknown" else 7, + ] + runner.finish_output(successful=True) kill.assert_not_called() assert runner.jobs == [] + stored = cloudai.models.output.Experiment.model_validate_json( + (runner.scenario_root / "experiment.json").read_text() + ) + assert stored.status == status + start = datetime.datetime(2026, 1, 2, 3, 4, tzinfo=datetime.timezone.utc) + expected = [] + for index, tr in enumerate(runs): + run_status = "completed" if index == 0 else status + metrics = [{"name": "Bandwidth", "value": 5 * (index + 1), "unit": "GB/s", "dimensions": []}] + if run_status != "completed": + metrics = [] + expected.append( + { + "id": tr.name, + "name": "sleep", + "description": "desc", + "status": run_status, + "path": str(runner.scenario_root / tr.name), + "metrics": metrics, + "dse": None, + "runs": [ + { + "path": str(tr.output_path.absolute()), + "jobid": "123", + "status": run_status, + "metrics": metrics, + "start": start + datetime.timedelta(seconds=1 if index == 0 else 5) + if run_status != "unknown" + else None, + "finish": start + datetime.timedelta(seconds=4 if index == 0 else 12) + if run_status != "unknown" + else None, + "duration": (3 if index == 0 else 7) if run_status != "unknown" else None, + "iteration": 0, + "step": 0, + } + ], + } + ) + assert [test.model_dump() for test in stored.tests] == expected def test_pre_test(nccl_tr: TestRun, sleep_tr: TestRun, slurm_system: SlurmSystem) -> None: @@ -624,7 +719,8 @@ def test_pre_test(nccl_tr: TestRun, sleep_tr: TestRun, slurm_system: SlurmSystem ) -def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None: +@pytest.mark.parametrize("failed_step", [None, 2]) +def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem, failed_step: int | None) -> None: tc = TestScenario(name="tc", test_runs=[dse_tr]) runner = SingleSbatchRunner(mode="run", system=slurm_system, test_scenario=tc, output_path=slurm_system.output_path) dse_tr.output_path = slurm_system.output_path / dse_tr.name @@ -632,7 +728,40 @@ def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None: trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.csv" trajectory_path.unlink(missing_ok=True) - runner.handle_dse() + dse_tr.test.agent_reward_function = "identity" + for tr in runner.all_trs: + runner.experiment_output.update_run( + dse_tr.name, + cloudai.models.output.Run( + path=str(tr.output_path), + jobid="123", + iteration=0, + step=tr.step, + status="failed" if tr.step == failed_step else "completed", + metrics=[cloudai.models.output.Metric(name="Bandwidth", value=tr.step * 12.5, unit="GB/s")], + ), + ) + with patch.object( + CloudAIGymEnv, + "get_observation", + side_effect=[{metric: float(step) for metric in dse_tr.test.agent_metrics} for step in (1, 2)], + ): + runner.handle_dse() + stored = cloudai.models.output.Experiment.model_validate_json( + (runner.scenario_root / "experiment.json").read_text() + ) + test = stored.tests[0] + best_step = 1 if failed_step else 2 + assert test.dse is not None + assert test.dse.model_dump() == { + "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, + "best_step": best_step, + "best_config": {"extra_env_vars.VAR1": f"value{best_step}"}, + } + assert [metric.model_dump() for metric in test.metrics] == [ + {"name": "Bandwidth", "value": best_step * 12.5, "unit": "GB/s", "dimensions": []}, + ] + assert len(test.runs) == 2 assert trajectory_path.exists() df = pd.read_csv(trajectory_path)