Add Slurm run and DSE results to experiment output

doc/reporting.rst

@@ -36,15 +36,26 @@ Unified experiment output
 
 CloudAI writes ``experiment.json`` in each scenario's results directory.
 
-The file contains scenario details, test cases, status, timing, and result paths. Standalone execution also records each
-run's process ID, status, iteration, timing, and workload metrics.
+The file contains scenario details and test cases under ``tests``. Standalone and Slurm execution records appear in each
+test case's ``runs`` list. Each record represents an iteration or DSE step and includes its number, process or Slurm job
+ID, status, timing, result path, and workload metrics.
 
-CloudAI updates the file when standalone runs start and finish, then finalizes it when scenario execution succeeds or
-fails. Timestamps use UTC; durations use seconds. Unknown timestamps are ``null``. Dry runs also produce scenario and
-test-case details without launching workloads.
+Timestamps use UTC; durations use seconds. Unknown timestamps are ``null``. A final status of ``unknown`` means the outcome
+could not be determined. Dry runs include scenario and test-case details without launching workloads.
 
-Metrics come from ``TestDefinition.metric_observations()``, independently of reporter settings. Each file update replaces
-the previous snapshot atomically. Metric extraction or write errors produce warnings without affecting execution.
+When a test case executes once successfully, ``tests[].metrics`` contains that execution's metrics. For DSE, it contains
+metrics from the successful step with the highest valid reward. The search space, selected step, and configuration appear
+in ``tests[].dse``. For example:
+
+.. code-block:: json
+
+   {
+     "space": {"extra_env_vars.NCCL_ALGO": ["Ring", "Tree"]},
+     "best_step": 2,
+     "best_config": {"extra_env_vars.NCCL_ALGO": "Tree"}
+   }
+
+Here, step 2 using ``Tree`` was selected. Its metrics appear in the test case's ``metrics`` list.
 
 
 .. _general-flow:

src/cloudai/_core/base_runner.py

@@ -113,7 +113,9 @@ def update_run_output(self, job: BaseJob, result: JobStatusResult | None = None)
         self.experiment_output.write()
 
     def finish_output(self, successful: bool) -> None:
-        status = "completed" if successful else "failed"
+        status: cloudai.models.output.Status = "completed" if successful else "failed"
+        if self.shutting_down and not any(test.status == "failed" for test in self.experiment_output.experiment.tests):
+            status = "cancelled"
         self.experiment_output.finish(status=status, finish=datetime.datetime.now(datetime.timezone.utc))
 
     def shutdown(self):

src/cloudai/cli/handlers.py

@@ -173,7 +173,11 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int:
             agent = agent_class(env, agent_config)
             logging.debug(f"Created agent {agent.__class__.__name__}.")
 
-            err |= agent.run()
+            env.update_output()
+            try:
+                err |= agent.run()
+            finally:
+                env.update_output()
     except Exception as exc:
         run_error = exc
         logging.exception("DSE job aborted by an unexpected error; generating reports before failing.")

src/cloudai/configurator/cloudai_gym.py

@@ -16,6 +16,7 @@
 
 import copy
 import logging
+import math
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
 
@@ -62,6 +63,34 @@ def __init__(
         self.trajectory = Trajectory(iteration_dir=self.iteration_dir)
         super().__init__()
 
+    def _ranked_dse_candidates(self) -> list[tuple[int, dict[str, str | int | float]]]:
+        """Return valid trial steps and configurations in descending reward order."""
+        tr = self.original_test_run
+        candidates = [
+            row
+            for row in self.trajectory.dataframe.to_dict(orient="records")
+            if math.isfinite(row["reward"])
+            and all(
+                math.isfinite(row[f"observation.{metric}"])
+                and row[f"observation.{metric}"] != self.rewards.metric_failure
+                for metric in tr.test.agent_metrics
+            )
+        ]
+        return [
+            (int(row["step"]), {key: row[f"action.{key}"] for key in tr.param_space})
+            for row in sorted(candidates, key=lambda row: row["reward"], reverse=True)
+        ]
+
+    def update_output(self) -> None:
+        """Publish the DSE recommendation without interrupting execution on output errors."""
+        try:
+            tr = self.original_test_run
+            output = self.runner.experiment_output
+            output.update_dse(str(tr.name), tr.param_space, self._ranked_dse_candidates())
+            output.write()
+        except Exception as exc:
+            logging.warning("Cannot update DSE output for %s: %s", self.original_test_run.name, exc)
+
     @property
     def upcoming_trial(self) -> int:
         """

src/cloudai/output.py

@@ -19,9 +19,23 @@
 import pathlib
 import tempfile
 
+import cloudai.metrics
 import cloudai.models.output
 
 
+def metric_output(observation: cloudai.metrics.MetricObservation) -> cloudai.models.output.Metric:
+    """Convert a canonical observation to an output metric."""
+    return cloudai.models.output.Metric(
+        name=observation.metric.display_name,
+        value=observation.value,
+        unit=observation.metric.unit,
+        dimensions=[
+            cloudai.models.output.Dimension(name=cloudai.metrics.dimension_label(key), value=str(value))
+            for key, value in sorted(observation.dimensions.items())
+        ],
+    )
+
+
 class ExperimentOutput:
     """Collect experiment results and publish snapshots."""
 
@@ -56,6 +70,25 @@ def update_test(self, test: cloudai.models.output.Test) -> None:
                 return
         self.experiment.tests.append(recorded)
 
+    def update_dse(
+        self,
+        test_id: str,
+        space: dict[str, list[str | int | float]],
+        candidates: list[tuple[int, dict[str, str | int | float]]],
+    ) -> None:
+        """Store the first ranked candidate with a successful run and its metrics."""
+        test = next((test for test in self.experiment.tests if test.id == test_id), None)
+        if test is None:
+            raise KeyError(f"Unknown experiment test: {test_id}")
+        completed_runs = {run.step: run for run in test.runs if run.status == "completed"}
+        for step, config in candidates:
+            if step in completed_runs:
+                test.dse = cloudai.models.output.DSE(space=space, best_step=step, best_config=config)
+                test.metrics = [metric.model_copy(deep=True) for metric in completed_runs[step].metrics]
+                return
+        test.dse = cloudai.models.output.DSE(space=space)
+        test.metrics = []
+
     def snapshot(self) -> cloudai.models.output.Experiment:
         """Return an independent snapshot without finalizing the experiment."""
         full = self.experiment.model_copy(deep=True)
@@ -87,15 +120,36 @@ def write(self) -> None:
                     logging.warning("Cannot remove temporary experiment output %s: %s", temporary_path, exc)
 
     def finish(self, status: cloudai.models.output.Status, finish: datetime.datetime | None) -> None:
+        for test in self.experiment.tests:
+            for run in test.runs:
+                if run.status in ("pending", "running"):
+                    run.status = "unknown"
+            self._update_test_status(test)
+            self._update_test_metrics(test)
+        if status == "completed":
+            statuses = {test.status for test in self.experiment.tests}
+            for outcome in ("failed", "cancelled", "unknown"):
+                if outcome in statuses:
+                    status = outcome
+                    break
+        for test in self.experiment.tests:
+            if test.status in ("pending", "running"):
+                test.status = "completed" if status == "completed" else "unknown"
         self.experiment.status = status
         self.experiment.finish = finish
-        if status == "completed":
-            for test in self.experiment.tests:
-                if test.status not in ("failed", "cancelled"):
-                    test.status = "completed"
         self._update_timing(self.experiment)
         self.write()
 
+    @staticmethod
+    def _update_test_metrics(test: cloudai.models.output.Test) -> None:
+        if test.dse is not None:
+            return
+        test.metrics = []
+        if len(test.runs) == 1:
+            run = test.runs[0]
+            if run.status == "completed" and run.step in (None, 0):
+                test.metrics = [metric.model_copy(deep=True) for metric in run.metrics]
+
     @staticmethod
     def _update_test_status(test: cloudai.models.output.Test) -> None:
         statuses = {run.status for run in test.runs}
@@ -105,6 +159,10 @@ def _update_test_status(test: cloudai.models.output.Test) -> None:
             test.status = "cancelled"
         elif "running" in statuses:
             test.status = "running"
+        elif "pending" in statuses:
+            test.status = "pending"
+        elif "unknown" in statuses:
+            test.status = "unknown"
         elif statuses == {"completed"}:
             test.status = "completed"
 

src/cloudai/systems/slurm/single_sbatch_runner.py

@@ -23,7 +23,7 @@
 
 from cloudai.configurator import CloudAIGymEnv
 from cloudai.configurator.env_params import EnvParams
-from cloudai.core import BaseJob, Registry, System, TestRun, TestScenario
+from cloudai.core import BaseJob, JobStatusResult, Registry, System, TestRun, TestScenario
 from cloudai.util import format_time_limit, parse_time_limit
 
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
@@ -244,6 +244,9 @@ def handle_dse(self):
     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 _submit_test(self, tr: TestRun) -> SlurmJob:
         with open(self.scenario_root / "cloudai_sbatch_script.sh", "w") as f:
             f.write(self.gen_sbatch_content())

src/cloudai/systems/slurm/slurm_job.py

@@ -18,9 +18,12 @@
 
 from cloudai.core import BaseJob
 
+from .slurm_metadata import SlurmJobMetadata
+
 
 @dataclass
 class SlurmJob(BaseJob):
     """A job class for execution on a Slurm system."""
 
     nodes: list[str] = field(default_factory=list, init=False)
+    metadata: SlurmJobMetadata | None = field(default=None, init=False)

src/cloudai/systems/slurm/slurm_runner.py

@@ -14,13 +14,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import logging
 from pathlib import Path
 from typing import cast
 
 import toml
 
-from cloudai.core import BaseJob, BaseRunner, System, TestRun, TestScenario
+import cloudai.models.output
+import cloudai.output
+from cloudai.core import BaseJob, BaseRunner, JobStatusResult, System, TestRun, TestScenario
 
 from .slurm_command_gen_strategy import SlurmCommandGenStrategy
 from .slurm_job import SlurmJob
@@ -36,6 +39,44 @@ def __init__(self, mode: str, system: System, test_scenario: TestScenario, outpu
         self.system = cast(SlurmSystem, system)
         self.pinned_nodes: dict[str, list[str]] = {}
 
+    def get_run_output(
+        self, job: BaseJob, tr: TestRun, result: JobStatusResult | None = None
+    ) -> cloudai.models.output.Run | None:
+        metadata = cast(SlurmJob, job).metadata
+        status: cloudai.models.output.Status = "pending"
+        metrics: list[cloudai.models.output.Metric] = []
+        if result is not None:
+            status = "completed" if result.is_successful else "failed"
+            if job.terminated_by_dependency or (metadata is not None and metadata.state.startswith("CANCELLED")):
+                status = "cancelled"
+            if status == "completed":
+                try:
+                    metrics = [
+                        cloudai.output.metric_output(observation)
+                        for observation in tr.test.metric_observations(self.system, tr)
+                    ]
+                except Exception as exc:
+                    logging.warning("Cannot extract output metrics for Slurm job %s: %s", job.id, exc)
+        return cloudai.models.output.Run(
+            path=str(tr.output_path.absolute()),
+            jobid=str(job.id),
+            status=status,
+            metrics=metrics,
+            start=self._output_timestamp(metadata.start_time) if metadata is not None else None,
+            finish=self._output_timestamp(metadata.end_time) if metadata is not None else None,
+            duration=metadata.elapsed_time_sec if metadata is not None else None,
+            iteration=tr.current_iteration,
+            step=tr.step,
+        )
+
+    @staticmethod
+    def _output_timestamp(value: str) -> datetime.datetime | None:
+        try:
+            timestamp = datetime.datetime.fromisoformat(value.replace("Z", "+00:00"))
+        except ValueError:
+            return None
+        return timestamp if timestamp.utcoffset() is not None else None
+
     def submit_test(self, tr: TestRun) -> None:
         if tr.pin_nodes and tr.name in self.pinned_nodes:
             tr.nodes = self.pinned_nodes[tr.name].copy()
@@ -119,7 +160,11 @@ def _get_job_metadata(
     def store_job_metadata(self, job: SlurmJob):
         system = cast(SlurmSystem, self.system)
         steps_metadata = [self._mock_job_metadata()] if self.mode == "dry-run" else system.get_job_status(job)
+        if not steps_metadata:
+            logging.warning("No Slurm accounting metadata available for job %s", job.id)
+            return
         slurm_job_file, job_meta = self._get_job_metadata(job, steps_metadata)
+        job.metadata = job_meta
 
         logging.debug(f"Storing job metadata for job {job.id} to {slurm_job_file}")
         with slurm_job_file.open("w") as job_file:

src/cloudai/systems/slurm/slurm_system.py

@@ -481,6 +481,7 @@ def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> list[SlurmSt
 
         retry_count = 0
         command = (
+            "TZ=UTC SLURM_TIME_FORMAT='%Y-%m-%dT%H:%M:%SZ' "
             f"sacct -j {job.id} --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine "
             "--delimiter='|' -p --noheader"
         )

src/cloudai/systems/standalone/standalone_runner.py

@@ -19,8 +19,8 @@
 from pathlib import Path
 from typing import cast
 
-import cloudai.metrics
 import cloudai.models.output
+import cloudai.output
 from cloudai.core import BaseJob, BaseRunner, JobIdRetrievalError, JobStatusResult, System, TestRun, TestScenario
 from cloudai.util import CommandShell
 
@@ -52,7 +52,7 @@ def get_run_output(
             if result.is_successful:
                 try:
                     observations = tr.test.metric_observations(self.system, tr)
-                    metrics = [self._metric_output(observation) for observation in observations]
+                    metrics = [cloudai.output.metric_output(observation) for observation in observations]
                 except Exception as exc:
                     logging.warning("Cannot extract output metrics for standalone job %s: %s", job.id, exc)
         return cloudai.models.output.Run(
@@ -70,22 +70,6 @@ def on_job_completion(self, job: BaseJob) -> None:
         standalone_job = cast(StandaloneJob, job)
         standalone_job.finish = datetime.datetime.now(datetime.timezone.utc)
 
-    @staticmethod
-    def _metric_output(observation: cloudai.metrics.MetricObservation) -> cloudai.models.output.Metric:
-        dimensions = [
-            cloudai.models.output.Dimension(
-                name=cloudai.metrics.dimension_label(key),
-                value=str(value),
-            )
-            for key, value in sorted(observation.dimensions.items())
-        ]
-        return cloudai.models.output.Metric(
-            name=observation.metric.display_name,
-            value=observation.value,
-            unit=observation.metric.unit,
-            dimensions=dimensions,
-        )
-
     def _submit_test(self, tr: TestRun) -> StandaloneJob:
         logging.info(f"Running test: {tr.name}")
         tr.output_path = self.get_job_output_path(tr)