From acf642ad60ce561ae262cb478bb56232ebf05a96 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 18 Sep 2026 17:42:33 +0200 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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 04/10] 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: From 2906ebb27e988a0a1d84ca21c912543083651d6a Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 18 Sep 2026 20:08:52 +0200 Subject: [PATCH 05/10] Add unified results for single-sbatch Slurm scenarios Signed-off-by: Ivan Podkidyshev --- .../systems/slurm/single_sbatch_runner.py | 84 ++++++++--- tests/test_single_sbatch_runner.py | 139 +++++++++++++++++- 2 files changed, 192 insertions(+), 31 deletions(-) diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 4aa0972de..78943933b 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/src/cloudai/systems/slurm/single_sbatch_runner.py @@ -21,6 +21,7 @@ from pathlib import Path from typing import Generator, Optional, cast +import cloudai.models.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 +206,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 +221,67 @@ 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) + if 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..5464453c3 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,108 @@ 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() + 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 +715,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 +724,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) From 5a424721699c5fcf20d50b45f6de8c8cdb1d559e Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 22 Sep 2026 13:53:18 +0200 Subject: [PATCH 06/10] Move shared handlers out of the CLI package Signed-off-by: Ivan Podkidyshev --- pyproject.toml | 12 +++++++++--- src/cloudai/cli/cli.py | 2 +- src/cloudai/{cli => }/handlers.py | 0 tests/test_acceptance.py | 2 +- tests/test_handlers.py | 16 ++++++++-------- tests/test_reporter.py | 2 +- 6 files changed, 20 insertions(+), 14 deletions(-) rename src/cloudai/{cli => }/handlers.py (100%) diff --git a/pyproject.toml b/pyproject.toml index 394878152..cb643ee67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -152,20 +152,26 @@ root_package = "cloudai" "cloudai.workloads", "cloudai.util", "cloudai.cli", + "cloudai.handlers", "cloudai.report_generator", ] [[tool.importlinter.contracts]] name = "Report generator is leaf dependency" type = "forbidden" - forbidden_modules = ["cloudai.workloads", "cloudai.cli"] - allow_indirect_imports = true # allow "from cloudai.core import ..." + forbidden_modules = ["cloudai.workloads", "cloudai.cli", "cloudai.handlers"] + allow_indirect_imports = true # allow "from cloudai.core import ..." source_modules = ["cloudai.report_generator"] [[tool.importlinter.contracts]] name = "Util modules are leaf dependencies" type = "forbidden" - forbidden_modules = ["cloudai.systems", "cloudai.workloads", "cloudai.cli"] + forbidden_modules = [ + "cloudai.systems", + "cloudai.workloads", + "cloudai.cli", + "cloudai.handlers", + ] allow_indirect_imports = true source_modules = ["cloudai.util"] diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 0aa8acdad..66d478742 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -21,7 +21,7 @@ import click -from .handlers import ( +from cloudai.handlers import ( handle_dry_run_and_run, handle_generate_report, handle_install_and_uninstall, diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/handlers.py similarity index 100% rename from src/cloudai/cli/handlers.py rename to src/cloudai/handlers.py diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 53105c5d4..c28d98f05 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -27,8 +27,8 @@ import cloudai.models.output from cloudai.cli import setup_logging -from cloudai.cli.handlers import handle_dry_run_and_run from cloudai.core import CommandGenStrategy, GitRepo, TestDefinition, TestRun, TestScenario +from cloudai.handlers import handle_dry_run_and_run from cloudai.models.scenario import TestRunDetails from cloudai.systems.slurm import SlurmCommandGenStrategy, SlurmRunner, SlurmSystem from cloudai.workloads.ai_dynamo import ( diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 18993cabd..5609024b9 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -25,14 +25,6 @@ from pydantic import Field import cloudai.models.output -from cloudai.cli.handlers import ( - handle_dse_job, - prepare_installation, - validate_domain_randomization_active, - verify_system_configs, - verify_test_configs, - verify_test_scenarios, -) from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParamSpec from cloudai.core import ( @@ -48,6 +40,14 @@ TestScenario, TestScenarioParsingError, ) +from cloudai.handlers import ( + handle_dse_job, + prepare_installation, + validate_domain_randomization_active, + verify_system_configs, + verify_test_configs, + verify_test_scenarios, +) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import CmdArgs, TestDefinition from cloudai.reporter import StatusReporter diff --git a/tests/test_reporter.py b/tests/test_reporter.py index 045eb50d6..b6e3a0f65 100644 --- a/tests/test_reporter.py +++ b/tests/test_reporter.py @@ -26,8 +26,8 @@ import toml from cloudai import TestRun, TestScenario -from cloudai.cli.handlers import generate_reports from cloudai.core import CommandGenStrategy, Registry, Reporter, System +from cloudai.handlers import generate_reports from cloudai.models.scenario import ReportConfig, TestRunDetails from cloudai.report_generator.dse_report import build_dse_summaries from cloudai.reporter import DSEReporter, JUnitReporter, PerTestReporter, ReportItem, StatusReporter, TarballReporter From ed5c9c6f1764b72a415db13b19f703592b9e7c55 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 22 Sep 2026 15:50:25 +0200 Subject: [PATCH 07/10] Restore deprecated CLI handler compatibility imports Signed-off-by: Ivan Podkidyshev --- doc/DEV.rst | 7 +++ src/cloudai/cli/handlers.py | 65 ++++++++++++++++++++++++++++ tests/cli/__init__.py | 15 +++++++ tests/cli/test_handlers.py | 86 +++++++++++++++++++++++++++++++++++++ 4 files changed, 173 insertions(+) create mode 100644 src/cloudai/cli/handlers.py create mode 100644 tests/cli/__init__.py create mode 100644 tests/cli/test_handlers.py diff --git a/doc/DEV.rst b/doc/DEV.rst index e2a6c27c6..c1ca82f19 100644 --- a/doc/DEV.rst +++ b/doc/DEV.rst @@ -31,6 +31,13 @@ This chapter targets developers who want to contribute to the project's core. runners --> _core systems --> _core +Handler imports +--------------- + +Shared orchestration lives in ``cloudai.handlers``. The CLI calls this module directly. +Existing imports from ``cloudai.cli.handlers`` continue to work and emit ``DeprecationWarning`` when called. +These handlers are internal implementation details; avoid calling them directly from integrations. + Core Modules ------------ diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py new file mode 100644 index 000000000..3f3d91ab1 --- /dev/null +++ b/src/cloudai/cli/handlers.py @@ -0,0 +1,65 @@ +# 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 functools +import typing +import warnings + +import cloudai.handlers + +_P = typing.ParamSpec("_P") +_R = typing.TypeVar("_R") + + +def _deprecated(function: typing.Callable[_P, _R]) -> typing.Callable[_P, _R]: + @functools.wraps(function) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + warnings.warn( + f"cloudai.cli.handlers.{function.__name__} is deprecated; do not call CLI handlers directly.", + DeprecationWarning, + stacklevel=2, + ) + return getattr(cloudai.handlers, function.__name__)(*args, **kwargs) + + return wrapper + + +_log_installation_dirs = _deprecated(cloudai.handlers._log_installation_dirs) +handle_install_and_uninstall = _deprecated(cloudai.handlers.handle_install_and_uninstall) +prepare_installation = _deprecated(cloudai.handlers.prepare_installation) +_scenario_installables = _deprecated(cloudai.handlers._scenario_installables) +handle_dse_job = _deprecated(cloudai.handlers.handle_dse_job) +_record_run_failure = _deprecated(cloudai.handlers._record_run_failure) +generate_reports = _deprecated(cloudai.handlers.generate_reports) +handle_non_dse_job = _deprecated(cloudai.handlers.handle_non_dse_job) +register_signal_handlers = _deprecated(cloudai.handlers.register_signal_handlers) +_setup_system_and_scenario = _deprecated(cloudai.handlers._setup_system_and_scenario) +_handle_single_sbatch = _deprecated(cloudai.handlers._handle_single_sbatch) +_check_installation = _deprecated(cloudai.handlers._check_installation) +handle_dry_run_and_run = _deprecated(cloudai.handlers.handle_dry_run_and_run) +handle_generate_report = _deprecated(cloudai.handlers.handle_generate_report) +expand_file_list = _deprecated(cloudai.handlers.expand_file_list) +_ensure_kube_config_exists = _deprecated(cloudai.handlers._ensure_kube_config_exists) +verify_system_configs = _deprecated(cloudai.handlers.verify_system_configs) +verify_test_configs = _deprecated(cloudai.handlers.verify_test_configs) +verify_test_scenarios = _deprecated(cloudai.handlers.verify_test_scenarios) +handle_verify_all_configs = _deprecated(cloudai.handlers.handle_verify_all_configs) +load_tomls_by_type = _deprecated(cloudai.handlers.load_tomls_by_type) +handle_list_registered_items = _deprecated(cloudai.handlers.handle_list_registered_items) +validate_domain_randomization_active = _deprecated(cloudai.handlers.validate_domain_randomization_active) +load_test_toml_file = _deprecated(cloudai.handlers.load_test_toml_file) +format_toml_decode_error = _deprecated(cloudai.handlers.format_toml_decode_error) +prepare_output_dir = _deprecated(cloudai.handlers.prepare_output_dir) diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py new file mode 100644 index 000000000..3ac11851d --- /dev/null +++ b/tests/cli/__init__.py @@ -0,0 +1,15 @@ +# 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. diff --git a/tests/cli/test_handlers.py b/tests/cli/test_handlers.py new file mode 100644 index 000000000..936321f76 --- /dev/null +++ b/tests/cli/test_handlers.py @@ -0,0 +1,86 @@ +# 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 inspect +from unittest.mock import Mock, call + +import pytest + +import cloudai.cli.handlers +import cloudai.handlers + + +@pytest.mark.parametrize( + "name", + [ + "_log_installation_dirs", + "handle_install_and_uninstall", + "prepare_installation", + "_scenario_installables", + "handle_dse_job", + "_record_run_failure", + "generate_reports", + "handle_non_dse_job", + "register_signal_handlers", + "_setup_system_and_scenario", + "_handle_single_sbatch", + "_check_installation", + "handle_dry_run_and_run", + "handle_generate_report", + "expand_file_list", + "_ensure_kube_config_exists", + "verify_system_configs", + "verify_test_configs", + "verify_test_scenarios", + "handle_verify_all_configs", + "load_tomls_by_type", + "handle_list_registered_items", + "validate_domain_randomization_active", + "load_test_toml_file", + "format_toml_decode_error", + "prepare_output_dir", + ], +) +def test_legacy_handler(name: str, monkeypatch: pytest.MonkeyPatch): + original = getattr(cloudai.handlers, name) + legacy = getattr(cloudai.cli.handlers, name) + signature = inspect.signature(original) + assert inspect.signature(legacy) == signature + assert legacy.__doc__ == original.__doc__ + + target = Mock() + monkeypatch.setattr(cloudai.handlers, name, target) + kwargs = {parameter: object() for parameter in signature.parameters} + args = tuple(kwargs.values()) + required_args = tuple( + value + for parameter, value in kwargs.items() + if signature.parameters[parameter].default is inspect.Parameter.empty + ) + warning = rf"cloudai\.cli\.handlers\.{name} is deprecated" + for positional, keywords in [(args, {}), ((), kwargs), (required_args, {})]: + with pytest.warns(DeprecationWarning, match=warning) as recorded: + assert legacy(*positional, **keywords) is target.return_value + assert len(recorded) == 1 + assert recorded[0].filename == __file__ + + assert target.call_args_list == [call(*args), call(**kwargs), call(*required_args)] + + error = RuntimeError("handler failure") + target.side_effect = error + with pytest.warns(DeprecationWarning, match=warning), pytest.raises(RuntimeError) as raised: + legacy(*required_args) + assert raised.value is error From 00eb42ece7e39e9ae0ca8e416e5e15be34986256 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 22 Sep 2026 15:59:06 +0200 Subject: [PATCH 08/10] Keep compatibility shims limited to public functions Signed-off-by: Ivan Podkidyshev --- doc/DEV.rst | 7 --- src/cloudai/cli/handlers.py | 7 --- tests/cli/__init__.py | 15 ------- tests/cli/test_handlers.py | 86 ------------------------------------- 4 files changed, 115 deletions(-) delete mode 100644 tests/cli/__init__.py delete mode 100644 tests/cli/test_handlers.py diff --git a/doc/DEV.rst b/doc/DEV.rst index c1ca82f19..e2a6c27c6 100644 --- a/doc/DEV.rst +++ b/doc/DEV.rst @@ -31,13 +31,6 @@ This chapter targets developers who want to contribute to the project's core. runners --> _core systems --> _core -Handler imports ---------------- - -Shared orchestration lives in ``cloudai.handlers``. The CLI calls this module directly. -Existing imports from ``cloudai.cli.handlers`` continue to work and emit ``DeprecationWarning`` when called. -These handlers are internal implementation details; avoid calling them directly from integrations. - Core Modules ------------ diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 3f3d91ab1..40ddb434d 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -37,22 +37,15 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: return wrapper -_log_installation_dirs = _deprecated(cloudai.handlers._log_installation_dirs) handle_install_and_uninstall = _deprecated(cloudai.handlers.handle_install_and_uninstall) prepare_installation = _deprecated(cloudai.handlers.prepare_installation) -_scenario_installables = _deprecated(cloudai.handlers._scenario_installables) handle_dse_job = _deprecated(cloudai.handlers.handle_dse_job) -_record_run_failure = _deprecated(cloudai.handlers._record_run_failure) generate_reports = _deprecated(cloudai.handlers.generate_reports) handle_non_dse_job = _deprecated(cloudai.handlers.handle_non_dse_job) register_signal_handlers = _deprecated(cloudai.handlers.register_signal_handlers) -_setup_system_and_scenario = _deprecated(cloudai.handlers._setup_system_and_scenario) -_handle_single_sbatch = _deprecated(cloudai.handlers._handle_single_sbatch) -_check_installation = _deprecated(cloudai.handlers._check_installation) handle_dry_run_and_run = _deprecated(cloudai.handlers.handle_dry_run_and_run) handle_generate_report = _deprecated(cloudai.handlers.handle_generate_report) expand_file_list = _deprecated(cloudai.handlers.expand_file_list) -_ensure_kube_config_exists = _deprecated(cloudai.handlers._ensure_kube_config_exists) verify_system_configs = _deprecated(cloudai.handlers.verify_system_configs) verify_test_configs = _deprecated(cloudai.handlers.verify_test_configs) verify_test_scenarios = _deprecated(cloudai.handlers.verify_test_scenarios) diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py deleted file mode 100644 index 3ac11851d..000000000 --- a/tests/cli/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -# 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. diff --git a/tests/cli/test_handlers.py b/tests/cli/test_handlers.py deleted file mode 100644 index 936321f76..000000000 --- a/tests/cli/test_handlers.py +++ /dev/null @@ -1,86 +0,0 @@ -# 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 inspect -from unittest.mock import Mock, call - -import pytest - -import cloudai.cli.handlers -import cloudai.handlers - - -@pytest.mark.parametrize( - "name", - [ - "_log_installation_dirs", - "handle_install_and_uninstall", - "prepare_installation", - "_scenario_installables", - "handle_dse_job", - "_record_run_failure", - "generate_reports", - "handle_non_dse_job", - "register_signal_handlers", - "_setup_system_and_scenario", - "_handle_single_sbatch", - "_check_installation", - "handle_dry_run_and_run", - "handle_generate_report", - "expand_file_list", - "_ensure_kube_config_exists", - "verify_system_configs", - "verify_test_configs", - "verify_test_scenarios", - "handle_verify_all_configs", - "load_tomls_by_type", - "handle_list_registered_items", - "validate_domain_randomization_active", - "load_test_toml_file", - "format_toml_decode_error", - "prepare_output_dir", - ], -) -def test_legacy_handler(name: str, monkeypatch: pytest.MonkeyPatch): - original = getattr(cloudai.handlers, name) - legacy = getattr(cloudai.cli.handlers, name) - signature = inspect.signature(original) - assert inspect.signature(legacy) == signature - assert legacy.__doc__ == original.__doc__ - - target = Mock() - monkeypatch.setattr(cloudai.handlers, name, target) - kwargs = {parameter: object() for parameter in signature.parameters} - args = tuple(kwargs.values()) - required_args = tuple( - value - for parameter, value in kwargs.items() - if signature.parameters[parameter].default is inspect.Parameter.empty - ) - warning = rf"cloudai\.cli\.handlers\.{name} is deprecated" - for positional, keywords in [(args, {}), ((), kwargs), (required_args, {})]: - with pytest.warns(DeprecationWarning, match=warning) as recorded: - assert legacy(*positional, **keywords) is target.return_value - assert len(recorded) == 1 - assert recorded[0].filename == __file__ - - assert target.call_args_list == [call(*args), call(**kwargs), call(*required_args)] - - error = RuntimeError("handler failure") - target.side_effect = error - with pytest.warns(DeprecationWarning, match=warning), pytest.raises(RuntimeError) as raised: - legacy(*required_args) - assert raised.value is error From 5ee4cc3fa6713a57bdc3176687e891897ee72adf Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 28 Sep 2026 18:16:17 +0200 Subject: [PATCH 09/10] Place CLI config verification and listing beside commands Signed-off-by: Ivan Podkidyshev --- src/cloudai/cli/cli.py | 298 +++++++++++++++++++++++++++++++++++- src/cloudai/cli/handlers.py | 29 ++-- src/cloudai/handlers.py | 251 +----------------------------- tests/test_handlers.py | 4 +- 4 files changed, 309 insertions(+), 273 deletions(-) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 66d478742..b20f8e999 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -17,17 +17,26 @@ import argparse import logging import logging.config +from contextlib import contextmanager from pathlib import Path +from typing import List +from unittest.mock import Mock import click +import toml +import yaml -from cloudai.handlers import ( - handle_dry_run_and_run, - handle_generate_report, - handle_install_and_uninstall, - handle_list_registered_items, - handle_verify_all_configs, +from cloudai.configurator.env_params import validate_domain_randomization_active +from cloudai.core import ( + Parser, + Registry, + System, + TestParser, ) +from cloudai.handlers import handle_dry_run_and_run, handle_generate_report, handle_install_and_uninstall +from cloudai.parser import HOOK_ROOT +from cloudai.test_parser import load_test_toml_file +from cloudai.toml_utils import format_toml_decode_error def setup_logging(log_file: str, log_level: str) -> None: @@ -83,6 +92,279 @@ def setup_logging(log_file: str, log_level: str) -> None: logging.config.dictConfig(LOGGING_CONFIG) +def expand_file_list(root: Path, glob: str = "*.toml") -> tuple[int, List[Path]]: + if not root.exists(): + logging.error(f"{root} does not exist.") + return (1, []) + + test_tomls = [root] + if root.is_dir(): + test_tomls = list(root.glob(glob)) + if not test_tomls: + logging.error(f"No TOMLs found in {root}") + return (1, []) + + return (0, test_tomls) + + +@contextmanager +def _ensure_kube_config_exists(system_toml_path: Path, content: str): + try: + config_dict = toml.loads(content) + except toml.TomlDecodeError as e: + logging.error(format_toml_decode_error(system_toml_path, e, "system config")) + raise + + kube_config_path_str = config_dict.get("kube_config_path") + kube_config_path = Path(kube_config_path_str) if kube_config_path_str else Path.home() / ".kube" / "config" + + created_file = False + created_dir = False + + if not kube_config_path.exists(): + logging.warning(f"Kube config file '{kube_config_path}' not found. Creating a dummy one.") + if not kube_config_path.parent.exists(): + kube_config_path.parent.mkdir(parents=True, exist_ok=True) + created_dir = True + + dummy_config = { + "apiVersion": "v1", + "kind": "Config", + "preferences": {}, + "clusters": [{"name": "dummy-cluster", "cluster": {"server": "https://dummy-server"}}], + "users": [{"name": "dummy-user", "user": {"token": "dummy-token"}}], + "contexts": [{"name": "dummy-context", "context": {"cluster": "dummy-cluster", "user": "dummy-user"}}], + "current-context": "dummy-context", + } + kube_config_path.write_text(yaml.dump(dummy_config)) + created_file = True + else: + logging.debug(f"Kube config '{kube_config_path}' already exists. Skipping creation.") + + try: + yield kube_config_path + finally: + if created_file: + try: + kube_config_path.unlink() + logging.debug(f"Deleted temporary kube config: {kube_config_path}") + except Exception as e: + logging.warning(f"Failed to remove temporary kube config '{kube_config_path}': {e}") + if created_dir: + try: + kube_config_path.parent.rmdir() + logging.debug(f"Deleted kube config directory: {kube_config_path.parent}") + except OSError: + pass + + +def verify_system_configs(system_tomls: List[Path]) -> int: + nfailed = 0 + + for system_toml in system_tomls: + logging.debug(f"Verifying System: {system_toml}...") + content = system_toml.read_text() + + if 'scheduler = "kubernetes"' in content: + try: + with _ensure_kube_config_exists(system_toml, content): + Parser.parse_system(system_toml) + except Exception as e: + logging.error(f"Failed to verify system config {system_toml}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + else: + try: + Parser.parse_system(system_toml) + except Exception as e: + logging.error(f"Failed to verify system config {system_toml}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + + if nfailed: + logging.error(f"{nfailed} out of {len(system_tomls)} system configurations have issues.") + else: + logging.info(f"Checked systems: {len(system_tomls)}, all passed") + + return nfailed + + +def verify_test_configs(test_tomls: List[Path]) -> int: + nfailed = 0 + tp = TestParser([], None) # type: ignore + for test_toml in test_tomls: + logging.debug(f"Verifying Test: {test_toml}...") + try: + with test_toml.open() as fh: + tp.current_file = test_toml + tp.load_test_definition(load_test_toml_file(fh, test_toml)) + except Exception as e: + logging.error(f"Failed to verify Test: {test_toml}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + + if nfailed: + logging.error(f"{nfailed} out of {len(test_tomls)} test configurations have issues.") + else: + logging.info(f"Checked tests: {len(test_tomls)}, all passed") + + return nfailed + + +def verify_test_scenarios( + scenario_tomls: List[Path], test_tomls: list[Path], hook_tomls: List[Path], hook_test_tomls: list[Path] +) -> int: + system = Mock(spec=System, sol={}) + nfailed = 0 + for scenario_file in scenario_tomls: + logging.debug(f"Verifying Test Scenario: {scenario_file}...") + try: + tests = Parser.parse_tests(test_tomls, system) + hook_tests = Parser.parse_tests(hook_test_tomls, system) + hooks = Parser.parse_hooks(hook_tomls, system, {t.name: t for t in hook_tests}) + scenario = Parser.parse_test_scenario(scenario_file, system, {t.name: t for t in tests}, hooks) + validate_domain_randomization_active(scenario) + except Exception as e: + logging.error(f"Failed to verify Test Scenario: {scenario_file}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + + if nfailed: + logging.error(f"{nfailed} out of {len(scenario_tomls)} test scenarios have issues.") + else: + logging.info(f"Checked scenarios: {len(scenario_tomls)}, all passed") + + return nfailed + + +def handle_verify_all_configs(args: argparse.Namespace) -> int: + root: Path = args.configs_dir + err, tomls = expand_file_list(root, glob="**/*.toml") + if err: + return err + + err, hook_tomls = expand_file_list(HOOK_ROOT, glob="**/*.toml") + tomls += hook_tomls + logging.info(f"Found {len(hook_tomls)} hook TOMLs (always verified)") + + files = load_tomls_by_type(tomls) + + test_tomls = files["test"] + if args.tests_dir: + test_tomls = list(args.tests_dir.glob("*.toml")) + elif files["scenario"]: + logging.warning( + "Test configuration directory not provided, using all found test TOMLs in the specified directory." + ) + + nfailed = 0 + total_checked = 0 + + if files["system"]: + nfailed += verify_system_configs(files["system"]) + total_checked += len(files["system"]) + if test_tomls: + nfailed += verify_test_configs(test_tomls) + total_checked += len(test_tomls) + if files["scenario"]: + nfailed += verify_test_scenarios(files["scenario"], test_tomls, files["hook"], files["hook_test"]) + total_checked += len(files["scenario"]) + if files["unknown"]: + for unknown_file in files["unknown"]: + logging.error( + f"Unknown configuration file '{unknown_file}': could not classify as system, test, scenario, or hook." + ) + nfailed += len(files["unknown"]) + total_checked += len(files["unknown"]) + + if nfailed: + logging.error(f"{nfailed} out of {total_checked} configuration files have issues.") + else: + logging.info(f"Checked {total_checked} configuration files, all passed") + + return nfailed + + +def load_tomls_by_type(tomls: List[Path]) -> dict[str, List[Path]]: + files: dict[str, List[Path]] = { + "system": [], + "test": [], + "scenario": [], + "hook_test": [], + "hook": [], + "unknown": [], + } + for toml_file in tomls: + content = toml_file.read_text() + + is_in_hook_root = False + try: + toml_file.relative_to(HOOK_ROOT) + is_in_hook_root = True + except ValueError: + pass + + if is_in_hook_root: + if "test" in toml_file.parts: + files["hook_test"].append(toml_file) + else: + files["hook"].append(toml_file) + continue + + try: + toml_content = toml.loads(content) + except toml.TomlDecodeError: + files["unknown"].append(toml_file) + continue + + if not isinstance(toml_content, dict): + files["unknown"].append(toml_file) + continue + + if "scheduler" in toml_content: + files["system"].append(toml_file) + elif "test_template_name =" in content and "[[Tests]]" not in content: + files["test"].append(toml_file) + elif "[[Tests]]" in content: + files["scenario"].append(toml_file) + else: + files["unknown"].append(toml_file) + + return files + + +def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: C901 + registry = Registry() + if item_type.lower() == "reports": + print("Available scenario reports:") + for idx, (name, report) in enumerate(sorted(registry.scenario_reports.items()), start=1): + string = f'{idx}. "{name}" {report.__name__}' + if verbose: + string += f" (config={registry.report_configs[name].model_dump_json(indent=None)})" + print(string) + elif item_type.lower() == "agents": + print("Available agents:") + for idx, name in enumerate(registry.agent_names(), start=1): + agent = registry.get_agent(name) + string = f'{idx}. "{name}" class={agent.__name__}' + if verbose: + string += f"{agent.__doc__}" + print(string) + elif item_type.lower() == "reward-functions": + print("Available reward functions:") + for idx, name in enumerate(registry.reward_function_names(), start=1): + reward_function = registry.get_reward_function(name) + callable_name = getattr(reward_function, "__name__", type(reward_function).__name__) + string = f'{idx}. "{name}" function={callable_name}' + if verbose: + documentation = getattr(reward_function, "__doc__", None) + if documentation: + string += f" {documentation}" + print(string) + + return 0 + + def common_options(f): f = click.option( "--system-config", @@ -271,9 +553,9 @@ def verify_configs(configs_dir: Path, tests_dir: Path): raise click.exceptions.Exit(1 if handle_verify_all_configs(args) else 0) -@main.command() +@main.command(name="list") @click.argument("type", type=click.Choice(["reports", "agents", "reward-functions"], case_sensitive=False)) @click.option("-v", "--verbose", is_flag=True, default=False, help="Verbose output.") -def list(type: str, verbose: bool): +def list_registered(type: str, verbose: bool): """List available in Registry items.""" handle_list_registered_items(type, verbose) diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 40ddb434d..8bd80dcf8 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -18,7 +18,12 @@ import typing import warnings +import cloudai.cli.cli +import cloudai.configurator.env_params import cloudai.handlers +import cloudai.test_parser +import cloudai.toml_utils +import cloudai.util _P = typing.ParamSpec("_P") _R = typing.TypeVar("_R") @@ -32,7 +37,7 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: DeprecationWarning, stacklevel=2, ) - return getattr(cloudai.handlers, function.__name__)(*args, **kwargs) + return function(*args, **kwargs) return wrapper @@ -45,14 +50,14 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: register_signal_handlers = _deprecated(cloudai.handlers.register_signal_handlers) handle_dry_run_and_run = _deprecated(cloudai.handlers.handle_dry_run_and_run) handle_generate_report = _deprecated(cloudai.handlers.handle_generate_report) -expand_file_list = _deprecated(cloudai.handlers.expand_file_list) -verify_system_configs = _deprecated(cloudai.handlers.verify_system_configs) -verify_test_configs = _deprecated(cloudai.handlers.verify_test_configs) -verify_test_scenarios = _deprecated(cloudai.handlers.verify_test_scenarios) -handle_verify_all_configs = _deprecated(cloudai.handlers.handle_verify_all_configs) -load_tomls_by_type = _deprecated(cloudai.handlers.load_tomls_by_type) -handle_list_registered_items = _deprecated(cloudai.handlers.handle_list_registered_items) -validate_domain_randomization_active = _deprecated(cloudai.handlers.validate_domain_randomization_active) -load_test_toml_file = _deprecated(cloudai.handlers.load_test_toml_file) -format_toml_decode_error = _deprecated(cloudai.handlers.format_toml_decode_error) -prepare_output_dir = _deprecated(cloudai.handlers.prepare_output_dir) +expand_file_list = _deprecated(cloudai.cli.cli.expand_file_list) +verify_system_configs = _deprecated(cloudai.cli.cli.verify_system_configs) +verify_test_configs = _deprecated(cloudai.cli.cli.verify_test_configs) +verify_test_scenarios = _deprecated(cloudai.cli.cli.verify_test_scenarios) +handle_verify_all_configs = _deprecated(cloudai.cli.cli.handle_verify_all_configs) +load_tomls_by_type = _deprecated(cloudai.cli.cli.load_tomls_by_type) +handle_list_registered_items = _deprecated(cloudai.cli.cli.handle_list_registered_items) +validate_domain_randomization_active = _deprecated(cloudai.configurator.env_params.validate_domain_randomization_active) +load_test_toml_file = _deprecated(cloudai.test_parser.load_test_toml_file) +format_toml_decode_error = _deprecated(cloudai.toml_utils.format_toml_decode_error) +prepare_output_dir = _deprecated(cloudai.util.prepare_output_dir) diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 09402d619..e851efc79 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -19,13 +19,8 @@ import logging import signal import traceback -from contextlib import contextmanager from pathlib import Path -from typing import Callable, List, Optional -from unittest.mock import Mock - -import toml -import yaml +from typing import Callable, Optional from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( @@ -38,7 +33,6 @@ Registry, Runner, System, - TestParser, TestScenario, TestScenarioParsingError, ) @@ -46,8 +40,6 @@ from cloudai.models.workload import TestDefinition from cloudai.parser import HOOK_ROOT from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem -from cloudai.test_parser import load_test_toml_file -from cloudai.toml_utils import format_toml_decode_error from cloudai.util import prepare_output_dir @@ -394,247 +386,6 @@ def handle_generate_report(args: argparse.Namespace) -> int: return 0 -def expand_file_list(root: Path, glob: str = "*.toml") -> tuple[int, List[Path]]: - if not root.exists(): - logging.error(f"{root} does not exist.") - return (1, []) - - test_tomls = [root] - if root.is_dir(): - test_tomls = list(root.glob(glob)) - if not test_tomls: - logging.error(f"No TOMLs found in {root}") - return (1, []) - - return (0, test_tomls) - - -@contextmanager -def _ensure_kube_config_exists(system_toml_path: Path, content: str): - try: - config_dict = toml.loads(content) - except toml.TomlDecodeError as e: - logging.error(format_toml_decode_error(system_toml_path, e, "system config")) - raise - - kube_config_path_str = config_dict.get("kube_config_path") - kube_config_path = Path(kube_config_path_str) if kube_config_path_str else Path.home() / ".kube" / "config" - - created_file = False - created_dir = False - - if not kube_config_path.exists(): - logging.warning(f"Kube config file '{kube_config_path}' not found. Creating a dummy one.") - if not kube_config_path.parent.exists(): - kube_config_path.parent.mkdir(parents=True, exist_ok=True) - created_dir = True - - dummy_config = { - "apiVersion": "v1", - "kind": "Config", - "preferences": {}, - "clusters": [{"name": "dummy-cluster", "cluster": {"server": "https://dummy-server"}}], - "users": [{"name": "dummy-user", "user": {"token": "dummy-token"}}], - "contexts": [{"name": "dummy-context", "context": {"cluster": "dummy-cluster", "user": "dummy-user"}}], - "current-context": "dummy-context", - } - kube_config_path.write_text(yaml.dump(dummy_config)) - created_file = True - else: - logging.debug(f"Kube config '{kube_config_path}' already exists. Skipping creation.") - - try: - yield kube_config_path - finally: - if created_file: - try: - kube_config_path.unlink() - logging.debug(f"Deleted temporary kube config: {kube_config_path}") - except Exception as e: - logging.warning(f"Failed to remove temporary kube config '{kube_config_path}': {e}") - if created_dir: - try: - kube_config_path.parent.rmdir() - logging.debug(f"Deleted kube config directory: {kube_config_path.parent}") - except OSError: - pass - - -def verify_system_configs(system_tomls: List[Path]) -> int: - nfailed = 0 - - for system_toml in system_tomls: - logging.debug(f"Verifying System: {system_toml}...") - content = system_toml.read_text() - - if 'scheduler = "kubernetes"' in content: - try: - with _ensure_kube_config_exists(system_toml, content): - Parser.parse_system(system_toml) - except Exception as e: - logging.error(f"Failed to verify system config {system_toml}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - else: - try: - Parser.parse_system(system_toml) - except Exception as e: - logging.error(f"Failed to verify system config {system_toml}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - - if nfailed: - logging.error(f"{nfailed} out of {len(system_tomls)} system configurations have issues.") - else: - logging.info(f"Checked systems: {len(system_tomls)}, all passed") - - return nfailed - - -def verify_test_configs(test_tomls: List[Path]) -> int: - nfailed = 0 - tp = TestParser([], None) # type: ignore - for test_toml in test_tomls: - logging.debug(f"Verifying Test: {test_toml}...") - try: - with test_toml.open() as fh: - tp.current_file = test_toml - tp.load_test_definition(load_test_toml_file(fh, test_toml)) - except Exception as e: - logging.error(f"Failed to verify Test: {test_toml}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - - if nfailed: - logging.error(f"{nfailed} out of {len(test_tomls)} test configurations have issues.") - else: - logging.info(f"Checked tests: {len(test_tomls)}, all passed") - - return nfailed - - -def verify_test_scenarios( - scenario_tomls: List[Path], test_tomls: list[Path], hook_tomls: List[Path], hook_test_tomls: list[Path] -) -> int: - system = Mock(spec=System, sol={}) - nfailed = 0 - for scenario_file in scenario_tomls: - logging.debug(f"Verifying Test Scenario: {scenario_file}...") - try: - tests = Parser.parse_tests(test_tomls, system) - hook_tests = Parser.parse_tests(hook_test_tomls, system) - hooks = Parser.parse_hooks(hook_tomls, system, {t.name: t for t in hook_tests}) - scenario = Parser.parse_test_scenario(scenario_file, system, {t.name: t for t in tests}, hooks) - validate_domain_randomization_active(scenario) - except Exception as e: - logging.error(f"Failed to verify Test Scenario: {scenario_file}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - - if nfailed: - logging.error(f"{nfailed} out of {len(scenario_tomls)} test scenarios have issues.") - else: - logging.info(f"Checked scenarios: {len(scenario_tomls)}, all passed") - - return nfailed - - -def handle_verify_all_configs(args: argparse.Namespace) -> int: - root: Path = args.configs_dir - err, tomls = expand_file_list(root, glob="**/*.toml") - if err: - return err - - err, hook_tomls = expand_file_list(HOOK_ROOT, glob="**/*.toml") - tomls += hook_tomls - logging.info(f"Found {len(hook_tomls)} hook TOMLs (always verified)") - - files = load_tomls_by_type(tomls) - - test_tomls = files["test"] - if args.tests_dir: - test_tomls = list(args.tests_dir.glob("*.toml")) - elif files["scenario"]: - logging.warning( - "Test configuration directory not provided, using all found test TOMLs in the specified directory." - ) - - nfailed = 0 - total_checked = 0 - - if files["system"]: - nfailed += verify_system_configs(files["system"]) - total_checked += len(files["system"]) - if test_tomls: - nfailed += verify_test_configs(test_tomls) - total_checked += len(test_tomls) - if files["scenario"]: - nfailed += verify_test_scenarios(files["scenario"], test_tomls, files["hook"], files["hook_test"]) - total_checked += len(files["scenario"]) - if files["unknown"]: - for unknown_file in files["unknown"]: - logging.error( - f"Unknown configuration file '{unknown_file}': could not classify as system, test, scenario, or hook." - ) - nfailed += len(files["unknown"]) - total_checked += len(files["unknown"]) - - if nfailed: - logging.error(f"{nfailed} out of {total_checked} configuration files have issues.") - else: - logging.info(f"Checked {total_checked} configuration files, all passed") - - return nfailed - - -def load_tomls_by_type(tomls: List[Path]) -> dict[str, List[Path]]: - files: dict[str, List[Path]] = { - "system": [], - "test": [], - "scenario": [], - "hook_test": [], - "hook": [], - "unknown": [], - } - for toml_file in tomls: - content = toml_file.read_text() - - is_in_hook_root = False - try: - toml_file.relative_to(HOOK_ROOT) - is_in_hook_root = True - except ValueError: - pass - - if is_in_hook_root: - if "test" in toml_file.parts: - files["hook_test"].append(toml_file) - else: - files["hook"].append(toml_file) - continue - - try: - toml_content = toml.loads(content) - except toml.TomlDecodeError: - files["unknown"].append(toml_file) - continue - - if not isinstance(toml_content, dict): - files["unknown"].append(toml_file) - continue - - if "scheduler" in toml_content: - files["system"].append(toml_file) - elif "test_template_name =" in content and "[[Tests]]" not in content: - files["test"].append(toml_file) - elif "[[Tests]]" in content: - files["scenario"].append(toml_file) - else: - files["unknown"].append(toml_file) - - return files - - def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: C901 registry = Registry() if item_type.lower() == "reports": diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 5609024b9..766330c61 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -25,6 +25,7 @@ from pydantic import Field import cloudai.models.output +from cloudai.cli.cli import verify_system_configs, verify_test_configs, verify_test_scenarios from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParamSpec from cloudai.core import ( @@ -44,9 +45,6 @@ handle_dse_job, prepare_installation, validate_domain_randomization_active, - verify_system_configs, - verify_test_configs, - verify_test_scenarios, ) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import CmdArgs, TestDefinition From 145432dd9273c7c055c003a438f3c2a3933b0437 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 28 Sep 2026 18:17:18 +0200 Subject: [PATCH 10/10] Place CLI execution adapters beside commands Signed-off-by: Ivan Podkidyshev --- src/cloudai/cli/cli.py | 205 +++++++++++++++++++++++++++++++++++- src/cloudai/cli/handlers.py | 13 +-- src/cloudai/handlers.py | 204 +---------------------------------- tests/test_acceptance.py | 2 +- tests/test_handlers.py | 3 +- 5 files changed, 214 insertions(+), 213 deletions(-) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index b20f8e999..b2c920a95 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -17,26 +17,35 @@ import argparse import logging import logging.config +import signal from contextlib import contextmanager from pathlib import Path -from typing import List +from typing import Callable, List from unittest.mock import Mock import click import toml import yaml +import cloudai.handlers from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( + InstallStatusResult, + MissingTestError, Parser, Registry, + Runner, System, + TestDefinition, TestParser, + TestScenario, + TestScenarioParsingError, ) -from cloudai.handlers import handle_dry_run_and_run, handle_generate_report, handle_install_and_uninstall from cloudai.parser import HOOK_ROOT +from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem from cloudai.test_parser import load_test_toml_file from cloudai.toml_utils import format_toml_decode_error +from cloudai.util import prepare_output_dir def setup_logging(log_file: str, log_level: str) -> None: @@ -92,6 +101,84 @@ def setup_logging(log_file: str, log_level: str) -> None: logging.config.dictConfig(LOGGING_CONFIG) +def _log_installation_dirs(prefix: str, system: System) -> None: + logging.info(f"{prefix} '{system.install_path.absolute()}'. HF cache is {system.hf_home_path.absolute()}.") + + +def handle_install_and_uninstall(args: argparse.Namespace) -> int: + """ + Manage the installation or uninstallation process for CloudAI. + + Based on user-specified mode, utilizing the Installer class. + + Args: + args (argparse.Namespace): The parsed command-line arguments. + """ + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) + system, tests, scenario = parser.parse(args.tests_dir, args.test_scenario) + + system.update() + logging.info(f"System Name: {system.name}") + logging.info(f"Scheduler: {system.scheduler}") + + installables, installer = cloudai.handlers.prepare_installation(system, tests, scenario) + + rc = 0 + if args.mode == "install": + all_installed = installer.is_installed(installables) + if all_installed: + _log_installation_dirs("CloudAI is already installed into", system) + else: + logging.info("Not all components are ready") + result = installer.install(installables) + if result.success: + _log_installation_dirs("CloudAI is successfully installed into", system) + else: + logging.error(result.message) + rc = 1 + + elif args.mode == "uninstall": + logging.info("Uninstalling test templates.") + result = installer.uninstall(installables) + if result.success: + logging.info("Uninstallation successful.") + else: + logging.error(result.message) + rc = 1 + + return rc + + +def register_signal_handlers(signal_handler: Callable) -> None: + """Register signal handlers for handling termination-related signals.""" + signals = [ + signal.SIGINT, + signal.SIGTERM, + signal.SIGHUP, + signal.SIGQUIT, + ] + for sig in signals: + signal.signal(sig, signal_handler) + + +def handle_generate_report(args: argparse.Namespace) -> int: + """ + Generate a report based on the existing configuration and test results. + + Args: + args (argparse.Namespace): The parsed command-line arguments. + """ + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) + system, _, test_scenario = parser.parse(args.tests_dir, args.test_scenario) + assert test_scenario is not None + + cloudai.handlers.generate_reports(system, test_scenario, args.result_dir) + + logging.debug("Report generation completed.") + + return 0 + + def expand_file_list(root: Path, glob: str = "*.toml") -> tuple[int, List[Path]]: if not root.exists(): logging.error(f"{root} does not exist.") @@ -365,6 +452,120 @@ def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: return 0 +def _setup_system_and_scenario( + args: argparse.Namespace, +) -> tuple[System, TestScenario, list[TestDefinition]] | None: + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) + try: + system, tests, test_scenario = parser.parse(args.tests_dir, args.test_scenario) + except MissingTestError as e: + logging.error(e.message) + return None + + assert test_scenario is not None + + if args.output_dir: + system.output_path = args.output_dir.absolute() + + if not prepare_output_dir(system.output_path): + return None + + if args.mode == "dry-run": + system.monitor_interval = 1 + system.update() + + return system, test_scenario, tests + + +def _handle_single_sbatch(args: argparse.Namespace, system: System) -> bool: + if not args.single_sbatch: + return True + + if not isinstance(system, SlurmSystem): + logging.error("Single sbatch is only supported for Slurm systems.") + return False + + Registry().update_runner("slurm", SingleSbatchRunner) + return True + + +def _check_installation( + args: argparse.Namespace, system: System, tests: list[TestDefinition], test_scenario: TestScenario +) -> InstallStatusResult: + logging.info("Checking if workloads components are installed.") + installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) + + if args.enable_cache_without_check: + result = installer.mark_as_installed(installables) + else: + result = installer.is_installed(installables) + + return result + + +def handle_dry_run_and_run(args: argparse.Namespace) -> int: + setup_result = _setup_system_and_scenario(args) + if setup_result is None: + return 1 + system, test_scenario, tests = setup_result + + try: + validate_domain_randomization_active(test_scenario) + except TestScenarioParsingError as e: + logging.error(str(e)) + return 1 + + if not _handle_single_sbatch(args, system): + return 1 + + logging.info(f"System Name: {system.name}") + logging.info(f"Scheduler: {system.scheduler}") + logging.info(f"Test Scenario Name: {test_scenario.name}") + + result = _check_installation(args, system, tests, test_scenario) + if args.mode == "run" and not result.success: + logging.info("Not all workloads components are installed. Installing...") + installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) + + result = installer.install(installables) + if result.success: + _log_installation_dirs("CloudAI is successfully installed into", system) + else: + logging.error("Failed to install workloads components.") + logging.error(result.message) + return 1 + elif args.mode == "dry-run": + # simulate installation for dry-run + installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) + result = installer.mark_as_installed(installables) + if not result.success: + logging.warning("Failed to mark workloads components as installed for dry-run.") + + logging.info(test_scenario.pretty_print()) + + runner = Runner(args.mode, system, test_scenario) + register_signal_handlers(runner.cancel_on_signal) + logging.info(f"Scenario results will be stored at: {runner.runner.scenario_root}") + + successful = False + try: + runner.runner.experiment_output.write() + has_dse = any(tr.is_dse_job for tr in test_scenario.test_runs) + if args.single_sbatch or not has_dse: # in this mode cases are unrolled using grid search + successful = cloudai.handlers.handle_non_dse_job(runner, args) + return 0 + + if all(tr.is_dse_job for tr in test_scenario.test_runs): + result = cloudai.handlers.handle_dse_job(runner, args) + successful = result == 0 + return result + + logging.error("Mixing DSE and non-DSE jobs is not allowed.") + return 1 + finally: + runner.runner.finish_output(successful) + + def common_options(f): f = click.option( "--system-config", diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 8bd80dcf8..7e27d172d 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -42,14 +42,12 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: return wrapper -handle_install_and_uninstall = _deprecated(cloudai.handlers.handle_install_and_uninstall) +handle_install_and_uninstall = _deprecated(cloudai.cli.cli.handle_install_and_uninstall) prepare_installation = _deprecated(cloudai.handlers.prepare_installation) -handle_dse_job = _deprecated(cloudai.handlers.handle_dse_job) generate_reports = _deprecated(cloudai.handlers.generate_reports) -handle_non_dse_job = _deprecated(cloudai.handlers.handle_non_dse_job) -register_signal_handlers = _deprecated(cloudai.handlers.register_signal_handlers) -handle_dry_run_and_run = _deprecated(cloudai.handlers.handle_dry_run_and_run) -handle_generate_report = _deprecated(cloudai.handlers.handle_generate_report) +register_signal_handlers = _deprecated(cloudai.cli.cli.register_signal_handlers) +handle_dry_run_and_run = _deprecated(cloudai.cli.cli.handle_dry_run_and_run) +handle_generate_report = _deprecated(cloudai.cli.cli.handle_generate_report) expand_file_list = _deprecated(cloudai.cli.cli.expand_file_list) verify_system_configs = _deprecated(cloudai.cli.cli.verify_system_configs) verify_test_configs = _deprecated(cloudai.cli.cli.verify_test_configs) @@ -61,3 +59,6 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: load_test_toml_file = _deprecated(cloudai.test_parser.load_test_toml_file) format_toml_decode_error = _deprecated(cloudai.toml_utils.format_toml_decode_error) prepare_output_dir = _deprecated(cloudai.util.prepare_output_dir) + +handle_dse_job = _deprecated(cloudai.handlers.handle_dse_job) +handle_non_dse_job = _deprecated(cloudai.handlers.handle_non_dse_job) diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index e851efc79..2563e6e86 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -17,78 +17,22 @@ import argparse import copy import logging -import signal import traceback from pathlib import Path -from typing import Callable, Optional +from typing import Optional -from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( BaseInstaller, CloudAIGymEnv, Installable, - InstallStatusResult, - MissingTestError, - Parser, Registry, Runner, System, TestScenario, - TestScenarioParsingError, ) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import TestDefinition -from cloudai.parser import HOOK_ROOT -from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem -from cloudai.util import prepare_output_dir - - -def _log_installation_dirs(prefix: str, system: System) -> None: - logging.info(f"{prefix} '{system.install_path.absolute()}'. HF cache is {system.hf_home_path.absolute()}.") - - -def handle_install_and_uninstall(args: argparse.Namespace) -> int: - """ - Manage the installation or uninstallation process for CloudAI. - - Based on user-specified mode, utilizing the Installer class. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - """ - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - system, tests, scenario = parser.parse(args.tests_dir, args.test_scenario) - - system.update() - logging.info(f"System Name: {system.name}") - logging.info(f"Scheduler: {system.scheduler}") - - installables, installer = prepare_installation(system, tests, scenario) - - rc = 0 - if args.mode == "install": - all_installed = installer.is_installed(installables) - if all_installed: - _log_installation_dirs("CloudAI is already installed into", system) - else: - logging.info("Not all components are ready") - result = installer.install(installables) - if result.success: - _log_installation_dirs("CloudAI is successfully installed into", system) - else: - logging.error(result.message) - rc = 1 - - elif args.mode == "uninstall": - logging.info("Uninstalling test templates.") - result = installer.uninstall(installables) - if result.success: - logging.info("Uninstallation successful.") - else: - logging.error(result.message) - rc = 1 - - return rc +from cloudai.systems.slurm import SlurmSystem def prepare_installation( @@ -242,150 +186,6 @@ def handle_non_dse_job(runner: Runner, args: argparse.Namespace) -> bool: return successful -def register_signal_handlers(signal_handler: Callable) -> None: - """Register signal handlers for handling termination-related signals.""" - signals = [ - signal.SIGINT, - signal.SIGTERM, - signal.SIGHUP, - signal.SIGQUIT, - ] - for sig in signals: - signal.signal(sig, signal_handler) - - -def _setup_system_and_scenario( - args: argparse.Namespace, -) -> tuple[System, TestScenario, list[TestDefinition]] | None: - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - try: - system, tests, test_scenario = parser.parse(args.tests_dir, args.test_scenario) - except MissingTestError as e: - logging.error(e.message) - return None - - assert test_scenario is not None - - if args.output_dir: - system.output_path = args.output_dir.absolute() - - if not prepare_output_dir(system.output_path): - return None - - if args.mode == "dry-run": - system.monitor_interval = 1 - system.update() - - return system, test_scenario, tests - - -def _handle_single_sbatch(args: argparse.Namespace, system: System) -> bool: - if not args.single_sbatch: - return True - - if not isinstance(system, SlurmSystem): - logging.error("Single sbatch is only supported for Slurm systems.") - return False - - Registry().update_runner("slurm", SingleSbatchRunner) - return True - - -def _check_installation( - args: argparse.Namespace, system: System, tests: list[TestDefinition], test_scenario: TestScenario -) -> InstallStatusResult: - logging.info("Checking if workloads components are installed.") - installables, installer = prepare_installation(system, tests, test_scenario) - - if args.enable_cache_without_check: - result = installer.mark_as_installed(installables) - else: - result = installer.is_installed(installables) - - return result - - -def handle_dry_run_and_run(args: argparse.Namespace) -> int: - setup_result = _setup_system_and_scenario(args) - if setup_result is None: - return 1 - system, test_scenario, tests = setup_result - - try: - validate_domain_randomization_active(test_scenario) - except TestScenarioParsingError as e: - logging.error(str(e)) - return 1 - - if not _handle_single_sbatch(args, system): - return 1 - - logging.info(f"System Name: {system.name}") - logging.info(f"Scheduler: {system.scheduler}") - logging.info(f"Test Scenario Name: {test_scenario.name}") - - result = _check_installation(args, system, tests, test_scenario) - if args.mode == "run" and not result.success: - logging.info("Not all workloads components are installed. Installing...") - installables, installer = prepare_installation(system, tests, test_scenario) - - result = installer.install(installables) - if result.success: - _log_installation_dirs("CloudAI is successfully installed into", system) - else: - logging.error("Failed to install workloads components.") - logging.error(result.message) - return 1 - elif args.mode == "dry-run": - # simulate installation for dry-run - installables, installer = prepare_installation(system, tests, test_scenario) - result = installer.mark_as_installed(installables) - if not result.success: - logging.warning("Failed to mark workloads components as installed for dry-run.") - - logging.info(test_scenario.pretty_print()) - - runner = Runner(args.mode, system, test_scenario) - register_signal_handlers(runner.cancel_on_signal) - logging.info(f"Scenario results will be stored at: {runner.runner.scenario_root}") - - successful = False - try: - runner.runner.experiment_output.write() - has_dse = any(tr.is_dse_job for tr in test_scenario.test_runs) - if args.single_sbatch or not has_dse: # in this mode cases are unrolled using grid search - successful = handle_non_dse_job(runner, args) - return 0 - - if all(tr.is_dse_job for tr in test_scenario.test_runs): - result = handle_dse_job(runner, args) - successful = result == 0 - return result - - logging.error("Mixing DSE and non-DSE jobs is not allowed.") - return 1 - finally: - runner.runner.finish_output(successful) - - -def handle_generate_report(args: argparse.Namespace) -> int: - """ - Generate a report based on the existing configuration and test results. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - """ - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - system, _, test_scenario = parser.parse(args.tests_dir, args.test_scenario) - assert test_scenario is not None - - generate_reports(system, test_scenario, args.result_dir) - - logging.debug("Report generation completed.") - - return 0 - - def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: C901 registry = Registry() if item_type.lower() == "reports": diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index c28d98f05..bf6202bb5 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -27,8 +27,8 @@ import cloudai.models.output from cloudai.cli import setup_logging +from cloudai.cli.cli import handle_dry_run_and_run from cloudai.core import CommandGenStrategy, GitRepo, TestDefinition, TestRun, TestScenario -from cloudai.handlers import handle_dry_run_and_run from cloudai.models.scenario import TestRunDetails from cloudai.systems.slurm import SlurmCommandGenStrategy, SlurmRunner, SlurmSystem from cloudai.workloads.ai_dynamo import ( diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 766330c61..9f401ec73 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -27,7 +27,7 @@ import cloudai.models.output from cloudai.cli.cli import verify_system_configs, verify_test_configs, verify_test_scenarios from cloudai.configurator import CloudAIGymEnv -from cloudai.configurator.env_params import EnvParamSpec +from cloudai.configurator.env_params import EnvParamSpec, validate_domain_randomization_active from cloudai.core import ( BaseAgent, BaseAgentConfig, @@ -44,7 +44,6 @@ from cloudai.handlers import ( handle_dse_job, prepare_installation, - validate_domain_randomization_active, ) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import CmdArgs, TestDefinition