Add unified results for single-sbatch Slurm scenarios

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,9 @@ def run(self):
         finally:
             self.jobs.remove(job)
 
-        self.handle_dse()
-
         self.on_job_completion(job)
-        self.update_run_output(job)
+        self.update_run_output(job, JobStatusResult(is_successful=is_completed))
+        self.handle_dse()
 
     def handle_dse(self):
         registry = Registry()
@@ -221,31 +222,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:

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,15 +719,49 @@ 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
     dse_tr.output_path.mkdir(parents=True, exist_ok=True)
 
     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)