From acf642ad60ce561ae262cb478bb56232ebf05a96 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 18 Sep 2026 17:42:33 +0200 Subject: [PATCH 1/4] Add Slurm run and DSE results to experiment output Signed-off-by: Ivan Podkidyshev --- src/cloudai/_core/base_runner.py | 4 +- src/cloudai/cli/handlers.py | 40 ++++- src/cloudai/output.py | 47 +++++- .../systems/slurm/single_sbatch_runner.py | 5 +- src/cloudai/systems/slurm/slurm_job.py | 3 + src/cloudai/systems/slurm/slurm_runner.py | 47 +++++- src/cloudai/systems/slurm/slurm_system.py | 1 + .../systems/standalone/standalone_runner.py | 20 +-- tests/systems/slurm/test_runner.py | 142 ++++++++++++++++++ tests/systems/slurm/test_system.py | 5 + tests/test_handlers.py | 65 +++++++- tests/test_output.py | 57 ++++++- 12 files changed, 403 insertions(+), 33 deletions(-) create mode 100644 tests/systems/slurm/test_runner.py diff --git a/src/cloudai/_core/base_runner.py b/src/cloudai/_core/base_runner.py index dbd8f70e7..681611c8b 100644 --- a/src/cloudai/_core/base_runner.py +++ b/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): diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 313cbbed2..19f057a67 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -17,6 +17,7 @@ import argparse import copy import logging +import math import signal import traceback from contextlib import contextmanager @@ -27,6 +28,7 @@ import toml import yaml +import cloudai.models.output from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( BaseInstaller, @@ -130,6 +132,38 @@ def _scenario_installables(scenario: TestScenario) -> list[Installable]: return installables +def _update_dse_output(env: CloudAIGymEnv) -> None: + try: + output = env.runner.experiment_output + tr = env.original_test_run + test = next(test for test in output.snapshot().tests if test.id == tr.name) + completed_runs = { + run.step: run for run in test.runs if run.status == "completed" and run.iteration == tr.current_iteration + } + candidates = [ + row + for row in env.trajectory.dataframe.to_dict(orient="records") + if row["step"] in completed_runs + and math.isfinite(row["reward"]) + and all( + math.isfinite(row[f"observation.{metric}"]) + and row[f"observation.{metric}"] != env.rewards.metric_failure + for metric in tr.test.agent_metrics + ) + ] + best = max(candidates, key=lambda row: row["reward"], default=None) + test.dse = cloudai.models.output.DSE( + space=tr.param_space, + best_step=int(best["step"]) if best is not None else None, + best_config={key: best[f"action.{key}"] for key in tr.param_space} if best is not None else None, + ) + test.metrics = completed_runs[best["step"]].metrics if best is not None else [] + output.update_test(test) + output.write() + except Exception as exc: + logging.warning("Cannot update DSE output for %s: %s", env.original_test_run.name, exc) + + def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: registry = Registry() @@ -173,7 +207,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() + _update_dse_output(env) + try: + err |= agent.run() + finally: + _update_dse_output(env) except Exception as exc: run_error = exc logging.exception("DSE job aborted by an unexpected error; generating reports before failing.") diff --git a/src/cloudai/output.py b/src/cloudai/output.py index c8a39fb5d..3f65aa99e 100644 --- a/src/cloudai/output.py +++ b/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.""" @@ -87,15 +101,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 +140,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" diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 5db16fc30..4aa0972de 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/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()) diff --git a/src/cloudai/systems/slurm/slurm_job.py b/src/cloudai/systems/slurm/slurm_job.py index c5633b765..2b903e375 100644 --- a/src/cloudai/systems/slurm/slurm_job.py +++ b/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) diff --git a/src/cloudai/systems/slurm/slurm_runner.py b/src/cloudai/systems/slurm/slurm_runner.py index 93a177962..4928cabb2 100644 --- a/src/cloudai/systems/slurm/slurm_runner.py +++ b/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: diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 7b3d31010..8b7fb4aec 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/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" ) diff --git a/src/cloudai/systems/standalone/standalone_runner.py b/src/cloudai/systems/standalone/standalone_runner.py index 4e03ef0a8..fe2d6b9db 100644 --- a/src/cloudai/systems/standalone/standalone_runner.py +++ b/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) diff --git a/tests/systems/slurm/test_runner.py b/tests/systems/slurm/test_runner.py new file mode 100644 index 000000000..e065d5a3b --- /dev/null +++ b/tests/systems/slurm/test_runner.py @@ -0,0 +1,142 @@ +# 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. + +import datetime +import pathlib +from unittest import mock + +import pytest + +import cloudai.core +import cloudai.metrics +from cloudai.systems.slurm import SlurmJob, SlurmRunner, SlurmSystem +from cloudai.systems.slurm.slurm_metadata import SlurmStepMetadata + + +@pytest.mark.parametrize( + "successful,state,status", + [(True, "FAILED", "completed"), (False, "COMPLETED", "failed"), (False, "CANCELLED", "cancelled")], +) +@pytest.mark.parametrize("step", [0, 2]) +def test_slurm_run_output( + tmp_path: pathlib.Path, + base_tr: cloudai.core.TestRun, + slurm_system: SlurmSystem, + successful: bool, + state: str, + status: str, + step: int, +) -> None: + runner = SlurmRunner("run", slurm_system, cloudai.core.TestScenario(name="scenario", test_runs=[base_tr]), tmp_path) + base_tr.output_path.mkdir(parents=True) + base_tr.step = step + job = SlurmJob(base_tr, id=123) + runner.update_run_output(job) + metadata = SlurmStepMetadata( + job_id=123, + step_id="", + name="job", + state=state, + exit_code="1:0", + elapsed_time_sec=3, + start_time="2026-01-02T03:04:05Z", + end_time="2026-01-02T03:04:08Z", + submit_line="sbatch run.sh", + ) + observation = cloudai.metrics.MetricObservation(cloudai.metrics.BANDWIDTH, 12.5, {"size_bytes": 1024}) + with ( + mock.patch.object(SlurmSystem, "get_job_status", return_value=[metadata]) as get_metadata, + mock.patch.object( + runner, + "get_cmd_gen_strategy", + return_value=mock.Mock(gen_srun_command=lambda: "srun cmd", generate_test_command=lambda: ["cmd"]), + ), + mock.patch.object( + cloudai.core.TestDefinition, + "was_run_successful", + return_value=cloudai.core.JobStatusResult(is_successful=successful), + ), + mock.patch.object( + cloudai.core.TestDefinition, "metric_observations", return_value=[observation] + ) as get_metrics, + ): + runner.store_job_metadata(job) + runner.update_run_output(job, runner.get_job_status(job)) + get_metadata.assert_called_once_with(job) + assert get_metrics.call_count == int(successful) + + base_tr.step = 3 + runner.shutting_down = status == "cancelled" + runner.finish_output(successful=True) + experiment = runner.experiment_output.snapshot() + assert experiment.status == status + test = experiment.tests[0] + assert test.metrics == (test.runs[0].metrics if successful and step == 0 else []) + start = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc) + assert [run.model_dump() for run in test.runs] == [ + { + "path": str(base_tr.output_path.absolute()), + "jobid": "123", + "status": status, + "metrics": [ + { + "name": "Bandwidth", + "value": 12.5, + "unit": "GB/s", + "dimensions": [{"name": "Size", "value": "1024", "unit": "", "is_x": False}], + } + ] + if successful + else [], + "start": start, + "finish": start + datetime.timedelta(seconds=3), + "duration": 3, + "iteration": 0, + "step": step, + } + ] + + +@pytest.mark.parametrize("timestamp", ["Unknown", "", "2026-01-02T03:04:05"]) +def test_slurm_output_unknown_timing_and_metric_failure( + tmp_path: pathlib.Path, + base_tr: cloudai.core.TestRun, + slurm_system: SlurmSystem, + timestamp: str, + caplog: pytest.LogCaptureFixture, +) -> None: + runner = SlurmRunner("run", slurm_system, cloudai.core.TestScenario(name="scenario", test_runs=[base_tr]), tmp_path) + job = SlurmJob(base_tr, id=123) + with mock.patch.object(SlurmSystem, "get_job_status", return_value=[]): + runner.store_job_metadata(job) + with mock.patch.object( + cloudai.core.TestDefinition, "metric_observations", side_effect=ValueError("broken metrics") + ): + run = runner.get_run_output(job, base_tr, cloudai.core.JobStatusResult(is_successful=True)) + assert run is not None + assert run.model_dump() == { + "path": str(base_tr.output_path.absolute()), + "jobid": "123", + "status": "completed", + "metrics": [], + "start": None, + "finish": None, + "duration": None, + "iteration": 0, + "step": 0, + } + assert runner._output_timestamp(timestamp) is None + assert "broken metrics" in caplog.text diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 561f742ef..2d694fd87 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -902,6 +902,11 @@ def test_get_job_status(slurm_system: SlurmSystem, stdout: str, stderr: str, exp slurm_system.get_job_status(job) else: assert slurm_system.get_job_status(job) == expected + slurm_system.cmd_shell.execute.assert_called_with( + "TZ=UTC SLURM_TIME_FORMAT='%Y-%m-%dT%H:%M:%SZ' " + "sacct -j 1 --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine " + "--delimiter='|' -p --noheader" + ) sacct_output = """2623913,job,COMPLETED,0:0,2025-05-09T01:34:52,2025-05-09T01:59:27,1475,sbatch sbatch_script.sh, diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 560b2f409..beefeea7b 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -17,13 +17,14 @@ import argparse import copy from pathlib import Path -from typing import Any, ClassVar, Iterator, Optional +from typing import Any, ClassVar, Iterator, Optional, cast from unittest.mock import MagicMock import pandas as pd import pytest from pydantic import Field +import cloudai.models.output from cloudai.cli.handlers import ( handle_dse_job, prepare_installation, @@ -32,6 +33,7 @@ verify_test_configs, verify_test_scenarios, ) +from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParamSpec from cloudai.core import ( BaseAgent, @@ -397,6 +399,67 @@ def test_handle_dse_job_invokes_agent_run( assert CustomRunStubAgent.run_calls == 1 +@pytest.mark.parametrize("failed_steps", [[], [2], [1, 2]]) +def test_dse_output_selects_completed_trial( + slurm_system: SlurmSystem, + dse_tr: TestRun, + custom_run_agent_name: str, + monkeypatch: pytest.MonkeyPatch, + failed_steps: list[int], +) -> None: + dse_tr.test.agent = custom_run_agent_name + runner = Runner("dry-run", slurm_system, TestScenario(name="scenario", test_runs=[dse_tr])) + + def run(agent: CustomRunStubAgent) -> int: + env = cast(CloudAIGymEnv, agent.env) + for step, value in enumerate(["value1", "value2"], start=1): + env.runner.experiment_output.update_run( + dse_tr.name, + cloudai.models.output.Run( + path=str(env.iteration_dir / str(step)), + jobid=str(step), + step=step, + iteration=0, + status="failed" if step in failed_steps else "completed", + metrics=[cloudai.models.output.Metric(name="Bandwidth", value=12.5 * step, unit="GB/s")], + ), + ) + env.trajectory.append( + step=step, + action={"extra_env_vars.VAR1": value}, + reward=step, + observation={metric: float(step) for metric in dse_tr.test.agent_metrics}, + env_params={}, + ) + if failed_steps: + raise RuntimeError("trial failed") + return 0 + + monkeypatch.setattr(CustomRunStubAgent, "run", run) + if failed_steps: + with pytest.raises(RuntimeError, match="trial failed"): + handle_dse_job(runner, argparse.Namespace(mode="dry-run")) + else: + assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 0 + + stored = cloudai.models.output.Experiment.model_validate_json( + (runner.runner.scenario_root / "experiment.json").read_text() + ) + best_step = next((step for step in [2, 1] if step not in failed_steps), None) + assert stored.tests[0].dse is not None + assert stored.tests[0].dse.model_dump() == { + "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, + "best_step": best_step, + "best_config": {"extra_env_vars.VAR1": f"value{best_step}"} if best_step is not None else None, + } + assert [metric.model_dump() for metric in stored.tests[0].metrics] == ( + [{"name": "Bandwidth", "value": 12.5 * best_step, "unit": "GB/s", "dimensions": []}] + if best_step is not None + else [] + ) + assert len(stored.tests[0].runs) == 2 + + def test_handle_dse_job_propagates_agent_run_nonzero_rc( slurm_system: SlurmSystem, dse_tr: TestRun, diff --git a/tests/test_output.py b/tests/test_output.py index 93622f5b7..65e74916f 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -17,11 +17,16 @@ import datetime import pathlib +import pytest + import cloudai.models.output import cloudai.output -def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathlib.Path) -> None: +@pytest.mark.parametrize("status", ["completed", "failed", "cancelled"]) +def test_experiment_output_preserves_runs_and_finalizes_failure( + tmp_path: pathlib.Path, status: cloudai.models.output.Status +) -> None: start = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc) experiment = cloudai.models.output.Experiment( id="experiment", @@ -29,7 +34,10 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli status="running", path=str(tmp_path), start=start, - tests=[cloudai.models.output.Test(id="case", name="workload", path=str(tmp_path / "case"))], + tests=[ + cloudai.models.output.Test(id=case, name="workload", path=str(tmp_path / case)) + for case in ("case", "interrupted", "not-started") + ], ) experiment_output = cloudai.output.ExperimentOutput(experiment, tmp_path) first_run = cloudai.models.output.Run( @@ -44,7 +52,7 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli second_run = cloudai.models.output.Run( path=str(tmp_path / "case" / "1"), jobid="102", - status="running", + status="pending", start=start + datetime.timedelta(seconds=2), iteration=1, step=0, @@ -52,17 +60,22 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli experiment_output.update_run("case", first_run) experiment_output.update_run("case", second_run) + assert experiment_output.snapshot().tests[0].status == "pending" second_run.status = "failed" second_run.finish = start + datetime.timedelta(seconds=4) experiment_output.update_run("case", second_run) - experiment_output.finish("failed", start + datetime.timedelta(seconds=5)) + experiment_output.update_run( + "interrupted", + cloudai.models.output.Run(path=str(tmp_path / "interrupted" / "0"), jobid="103", status="pending"), + ) + experiment_output.finish(status, start + datetime.timedelta(seconds=5)) stored = cloudai.models.output.Experiment.model_validate_json((tmp_path / "experiment.json").read_text()) assert stored.model_dump() == { "id": "experiment", "name": "scenario", "description": None, - "status": "failed", + "status": "cancelled" if status == "cancelled" else "failed", "path": str(tmp_path), "start": start, "finish": start + datetime.timedelta(seconds=5), @@ -100,6 +113,38 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli }, ], "dse": None, - } + }, + { + "id": "interrupted", + "name": "workload", + "description": None, + "status": "unknown", + "path": str(tmp_path / "interrupted"), + "metrics": [], + "runs": [ + { + "path": str(tmp_path / "interrupted" / "0"), + "jobid": "103", + "status": "unknown", + "metrics": [], + "start": None, + "finish": None, + "duration": None, + "iteration": None, + "step": None, + } + ], + "dse": None, + }, + { + "id": "not-started", + "name": "workload", + "description": None, + "status": "unknown", + "path": str(tmp_path / "not-started"), + "metrics": [], + "runs": [], + "dse": None, + }, ], } From 83ac3fd90ee8e9726d188becdaab1b1f1caed482 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 21 Sep 2026 14:52:18 +0200 Subject: [PATCH 2/4] Separate DSE candidate ranking from experiment output Signed-off-by: Ivan Podkidyshev --- src/cloudai/cli/handlers.py | 38 ++--------------------- src/cloudai/configurator/cloudai_gym.py | 29 +++++++++++++++++ src/cloudai/output.py | 22 +++++++++++++ tests/test_handlers.py | 41 ++++++++++++++++++++++--- 4 files changed, 89 insertions(+), 41 deletions(-) diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 19f057a67..09402d619 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -17,7 +17,6 @@ import argparse import copy import logging -import math import signal import traceback from contextlib import contextmanager @@ -28,7 +27,6 @@ import toml import yaml -import cloudai.models.output from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( BaseInstaller, @@ -132,38 +130,6 @@ def _scenario_installables(scenario: TestScenario) -> list[Installable]: return installables -def _update_dse_output(env: CloudAIGymEnv) -> None: - try: - output = env.runner.experiment_output - tr = env.original_test_run - test = next(test for test in output.snapshot().tests if test.id == tr.name) - completed_runs = { - run.step: run for run in test.runs if run.status == "completed" and run.iteration == tr.current_iteration - } - candidates = [ - row - for row in env.trajectory.dataframe.to_dict(orient="records") - if row["step"] in completed_runs - and math.isfinite(row["reward"]) - and all( - math.isfinite(row[f"observation.{metric}"]) - and row[f"observation.{metric}"] != env.rewards.metric_failure - for metric in tr.test.agent_metrics - ) - ] - best = max(candidates, key=lambda row: row["reward"], default=None) - test.dse = cloudai.models.output.DSE( - space=tr.param_space, - best_step=int(best["step"]) if best is not None else None, - best_config={key: best[f"action.{key}"] for key in tr.param_space} if best is not None else None, - ) - test.metrics = completed_runs[best["step"]].metrics if best is not None else [] - output.update_test(test) - output.write() - except Exception as exc: - logging.warning("Cannot update DSE output for %s: %s", env.original_test_run.name, exc) - - def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: registry = Registry() @@ -207,11 +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__}.") - _update_dse_output(env) + env.update_output() try: err |= agent.run() finally: - _update_dse_output(env) + env.update_output() except Exception as exc: run_error = exc logging.exception("DSE job aborted by an unexpected error; generating reports before failing.") diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py index cdafd5ce5..0633871e9 100644 --- a/src/cloudai/configurator/cloudai_gym.py +++ b/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.current_iteration, 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: """ diff --git a/src/cloudai/output.py b/src/cloudai/output.py index 3f65aa99e..30cd9f5db 100644 --- a/src/cloudai/output.py +++ b/src/cloudai/output.py @@ -70,6 +70,28 @@ def update_test(self, test: cloudai.models.output.Test) -> None: return self.experiment.tests.append(recorded) + def update_dse( + self, + test_id: str, + iteration: int, + 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" and run.iteration == iteration + } + 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) diff --git a/tests/test_handlers.py b/tests/test_handlers.py index beefeea7b..511f3b627 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -399,19 +399,40 @@ def test_handle_dse_job_invokes_agent_run( assert CustomRunStubAgent.run_calls == 1 -@pytest.mark.parametrize("failed_steps", [[], [2], [1, 2]]) +@pytest.mark.parametrize( + "failed_steps,last_reward,last_observation,best_step", + [ + ([], 2.0, 2.0, 2), + ([2], 2.0, 2.0, 1), + ([1, 2], 2.0, 2.0, None), + ([], float("inf"), 2.0, 1), + ([], 2.0, float("nan"), 1), + ([], 2.0, -1.0, 1), + ([], 1.0, 2.0, 1), + ], +) def test_dse_output_selects_completed_trial( slurm_system: SlurmSystem, dse_tr: TestRun, custom_run_agent_name: str, monkeypatch: pytest.MonkeyPatch, failed_steps: list[int], + last_reward: float, + last_observation: float, + best_step: int | None, ) -> None: dse_tr.test.agent = custom_run_agent_name runner = Runner("dry-run", slurm_system, TestScenario(name="scenario", test_runs=[dse_tr])) def run(agent: CustomRunStubAgent) -> int: env = cast(CloudAIGymEnv, agent.env) + test = env.runner.experiment_output.snapshot().tests[0] + assert test.dse is not None + assert test.dse.model_dump() == { + "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, + "best_step": None, + "best_config": None, + } for step, value in enumerate(["value1", "value2"], start=1): env.runner.experiment_output.update_run( dse_tr.name, @@ -427,10 +448,21 @@ def run(agent: CustomRunStubAgent) -> int: env.trajectory.append( step=step, action={"extra_env_vars.VAR1": value}, - reward=step, - observation={metric: float(step) for metric in dse_tr.test.agent_metrics}, + reward=last_reward if step == 2 else 1.0, + observation={metric: last_observation if step == 2 else 1.0 for metric in dse_tr.test.agent_metrics}, env_params={}, ) + env.runner.experiment_output.update_run( + dse_tr.name, + cloudai.models.output.Run( + path=str(env.iteration_dir.parent / "1" / "2"), + jobid="other-iteration", + step=2, + iteration=1, + status="completed", + metrics=[cloudai.models.output.Metric(name="Bandwidth", value=100, unit="GB/s")], + ), + ) if failed_steps: raise RuntimeError("trial failed") return 0 @@ -445,7 +477,6 @@ def run(agent: CustomRunStubAgent) -> int: stored = cloudai.models.output.Experiment.model_validate_json( (runner.runner.scenario_root / "experiment.json").read_text() ) - best_step = next((step for step in [2, 1] if step not in failed_steps), None) assert stored.tests[0].dse is not None assert stored.tests[0].dse.model_dump() == { "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, @@ -457,7 +488,7 @@ def run(agent: CustomRunStubAgent) -> int: if best_step is not None else [] ) - assert len(stored.tests[0].runs) == 2 + assert len(stored.tests[0].runs) == 3 def test_handle_dse_job_propagates_agent_run_nonzero_rc( From 6acd5c9455fa241cfbf8b3c5d461804c205d088d Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 21 Sep 2026 15:21:29 +0200 Subject: [PATCH 3/4] Match DSE output runs by step without iteration filtering Signed-off-by: Ivan Podkidyshev --- src/cloudai/configurator/cloudai_gym.py | 2 +- src/cloudai/output.py | 5 +---- tests/test_handlers.py | 14 +------------- 3 files changed, 3 insertions(+), 18 deletions(-) diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py index 0633871e9..a472da9b5 100644 --- a/src/cloudai/configurator/cloudai_gym.py +++ b/src/cloudai/configurator/cloudai_gym.py @@ -86,7 +86,7 @@ def update_output(self) -> None: try: tr = self.original_test_run output = self.runner.experiment_output - output.update_dse(str(tr.name), tr.current_iteration, tr.param_space, self._ranked_dse_candidates()) + 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) diff --git a/src/cloudai/output.py b/src/cloudai/output.py index 30cd9f5db..ada9aff34 100644 --- a/src/cloudai/output.py +++ b/src/cloudai/output.py @@ -73,7 +73,6 @@ def update_test(self, test: cloudai.models.output.Test) -> None: def update_dse( self, test_id: str, - iteration: int, space: dict[str, list[str | int | float]], candidates: list[tuple[int, dict[str, str | int | float]]], ) -> None: @@ -81,9 +80,7 @@ def update_dse( 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" and run.iteration == iteration - } + 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) diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 511f3b627..18993cabd 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -440,7 +440,6 @@ def run(agent: CustomRunStubAgent) -> int: path=str(env.iteration_dir / str(step)), jobid=str(step), step=step, - iteration=0, status="failed" if step in failed_steps else "completed", metrics=[cloudai.models.output.Metric(name="Bandwidth", value=12.5 * step, unit="GB/s")], ), @@ -452,17 +451,6 @@ def run(agent: CustomRunStubAgent) -> int: observation={metric: last_observation if step == 2 else 1.0 for metric in dse_tr.test.agent_metrics}, env_params={}, ) - env.runner.experiment_output.update_run( - dse_tr.name, - cloudai.models.output.Run( - path=str(env.iteration_dir.parent / "1" / "2"), - jobid="other-iteration", - step=2, - iteration=1, - status="completed", - metrics=[cloudai.models.output.Metric(name="Bandwidth", value=100, unit="GB/s")], - ), - ) if failed_steps: raise RuntimeError("trial failed") return 0 @@ -488,7 +476,7 @@ def run(agent: CustomRunStubAgent) -> int: if best_step is not None else [] ) - assert len(stored.tests[0].runs) == 3 + assert len(stored.tests[0].runs) == 2 def test_handle_dse_job_propagates_agent_run_nonzero_rc( From bb70189798a1c0f57920b050d2eb785c860a9ac4 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 21 Sep 2026 16:08:48 +0200 Subject: [PATCH 4/4] Document Slurm experiment output and DSE results Signed-off-by: Ivan Podkidyshev --- doc/reporting.rst | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/doc/reporting.rst b/doc/reporting.rst index 87034e7ff..e55cb4d25 100644 --- a/doc/reporting.rst +++ b/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: