From e60e47dfbab483d06138e94ac07127f8677441e2 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Wed, 30 Sep 2026 13:27:14 +0200 Subject: [PATCH 01/13] Add path-based experiment execution API Signed-off-by: Ivan Podkidyshev --- doc/api.rst | 29 +++++ doc/index.rst | 2 + src/cloudai/_core/runner.py | 36 +++++- src/cloudai/api.py | 60 ++++++++++ src/cloudai/cli/cli.py | 53 ++++++++- src/cloudai/handlers.py | 220 +++++++++++++++++------------------- src/cloudai/parser.py | 52 +++++---- tests/test_acceptance.py | 4 +- tests/test_api.py | 87 ++++++++++++++ 9 files changed, 392 insertions(+), 151 deletions(-) create mode 100644 doc/api.rst create mode 100644 src/cloudai/api.py create mode 100644 tests/test_api.py diff --git a/doc/api.rst b/doc/api.rst new file mode 100644 index 000000000..8f6c5c3de --- /dev/null +++ b/doc/api.rst @@ -0,0 +1,29 @@ +Python API +========== + +``cloudai.api`` lets a Python program run a scenario and find its saved results. +Pass configuration file paths as ``pathlib.Path`` objects: + +.. code-block:: python + + from pathlib import Path + + import cloudai.api + + system = Path("conf/common/system/example_slurm_cluster.toml") + scenario = Path("conf/common/test_scenario/sleep.toml") + + experiment = cloudai.api.run_experiment( + scenario, system, tests_dir=Path("conf/common/test") + ) + print(experiment.status, experiment.path) + +``run_experiment`` waits for the scenario to finish and returns an ``Experiment`` +model with the same data saved in ``experiment.json``. Workload failures appear in +its status. Configuration and setup errors raise exceptions. Pass ``hook_dir`` to +use a custom hook directory. + +``list_experiments(system)`` returns ``(experiment_id, result_directory)`` pairs +from the system's configured output directory. Directories without +``experiment.json`` are skipped. A missing output directory gives an empty list; +an invalid experiment file raises an error. diff --git a/doc/index.rst b/doc/index.rst index 48fb1332d..f9ace1234 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -13,6 +13,7 @@ This document contains the following chapters: Tutorial USER_GUIDE + api reporting systems workloads/index @@ -21,6 +22,7 @@ This document contains the following chapters: - :doc:`Tutorial` - :doc:`USER_GUIDE` +- :doc:`api` - :doc:`reporting` - :doc:`systems` - :doc:`workloads/index` diff --git a/src/cloudai/_core/runner.py b/src/cloudai/_core/runner.py index ee636dd30..17552c6bd 100644 --- a/src/cloudai/_core/runner.py +++ b/src/cloudai/_core/runner.py @@ -16,6 +16,7 @@ import datetime import logging +from pathlib import Path from types import FrameType from typing import Optional @@ -44,11 +45,32 @@ class Runner: test_scenario (TestScenario): The test scenario to be executed. """ - def __init__(self, mode: str, system: System, test_scenario: TestScenario): + def __init__( + self, + mode: str, + system: System, + test_scenario: TestScenario, + *, + runner_class: type[BaseRunner] | None = None, + output_path: Path | None = None, + ): logging.info(f"Initializing Runner [{mode.upper()}] mode") - self.runner = self.create_runner(mode, system, test_scenario) - - def create_runner(self, mode: str, system: System, test_scenario: TestScenario) -> BaseRunner: + if runner_class is None and output_path is None: + self.runner = self.create_runner(mode, system, test_scenario) + else: + self.runner = self.create_runner( + mode, system, test_scenario, runner_class=runner_class, output_path=output_path + ) + + def create_runner( + self, + mode: str, + system: System, + test_scenario: TestScenario, + *, + runner_class: type[BaseRunner] | None = None, + output_path: Path | None = None, + ) -> BaseRunner: """ Dynamically create a runner instance based on the system's scheduler type. @@ -56,6 +78,8 @@ def create_runner(self, mode: str, system: System, test_scenario: TestScenario) mode (str): The operation mode ('dry-run', 'run'). system (System): The system configuration. test_scenario (TestScenario): The test scenario to run. + runner_class (type[BaseRunner] | None): Override for this invocation. + output_path (Path | None): Exact scenario result directory, if allocated. Returns: BaseRunner: A runner instance suitable for the system. @@ -69,13 +93,13 @@ def create_runner(self, mode: str, system: System, test_scenario: TestScenario) msg = f"No runner registered for scheduler: {scheduler_type}" logging.error(msg) raise NotImplementedError(msg) - runner_class = registry.runners_map[scheduler_type] + runner_class = runner_class or registry.runners_map[scheduler_type] logging.info(f"Creating {runner_class.__name__}") if not system.output_path.exists(): system.output_path.mkdir() current_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - results_root = system.output_path / f"{test_scenario.name}_{current_time}" + results_root = output_path or system.output_path / f"{test_scenario.name}_{current_time}" return runner_class(mode, system, test_scenario, results_root) diff --git a/src/cloudai/api.py b/src/cloudai/api.py new file mode 100644 index 000000000..b5dca98ce --- /dev/null +++ b/src/cloudai/api.py @@ -0,0 +1,60 @@ +# 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 pathlib + +import cloudai.handlers +import cloudai.models.output +from cloudai.core import Parser + + +def run_experiment( + scenario: pathlib.Path, + system: pathlib.Path, + *, + tests_dir: pathlib.Path | None = None, + hook_dir: pathlib.Path | None = None, +) -> cloudai.models.output.Experiment: + """ + Run a scenario synchronously and return its experiment snapshot. + + Configuration and setup errors raise exceptions. Workload failures are + reflected in the returned experiment status. + """ + parsed_system, tests, parsed_scenario = cloudai.handlers.load_experiment( + scenario.expanduser().resolve(), + system.expanduser().resolve(), + tests_dir=tests_dir, + hook_dir=hook_dir, + ) + runner = cloudai.handlers.create_experiment_runner(parsed_system, parsed_scenario, unique_result_dir=True) + cloudai.handlers.execute_experiment(runner, tests) + return runner.runner.experiment_output.snapshot() + + +def list_experiments(system: pathlib.Path) -> list[tuple[str, pathlib.Path]]: + """List saved experiment IDs and result directories for a system.""" + output_dir = Parser.parse_system(system.expanduser().resolve()).output_path.expanduser().resolve() + if not output_dir.exists(): + return [] + if not output_dir.is_dir(): + raise NotADirectoryError(output_dir) + + experiments = [] + for path in sorted(output_dir.glob("*/experiment.json")): + experiment = cloudai.models.output.Experiment.model_validate_json(path.read_text(encoding="utf-8")) + experiments.append((experiment.id, path.parent)) + return experiments diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 66d478742..93b5488f6 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -17,12 +17,15 @@ import argparse import logging import logging.config +import signal from pathlib import Path +from typing import Callable import click +import cloudai.core +import cloudai.handlers from cloudai.handlers import ( - handle_dry_run_and_run, handle_generate_report, handle_install_and_uninstall, handle_list_registered_items, @@ -30,6 +33,50 @@ ) +def register_signal_handlers(signal_handler: Callable) -> None: + """Register the CLI's termination signals for a running experiment.""" + for sig in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP, signal.SIGQUIT): + signal.signal(sig, signal_handler) + + +def _handle_dry_run_and_run(args: argparse.Namespace) -> int: + """Translate CLI options and experiment outcomes into a process exit code.""" + try: + system, tests, scenario = cloudai.handlers.load_experiment( + args.test_scenario, args.system_config, tests_dir=args.tests_dir, hook_dir=args.hook_dir + ) + runner = cloudai.handlers.create_experiment_runner( + system, + scenario, + mode=args.mode, + output_dir=args.output_dir, + single_sbatch=args.single_sbatch, + ) + except ( + cloudai.core.MissingTestError, + cloudai.core.SystemConfigParsingError, + cloudai.core.TestConfigParsingError, + cloudai.core.TestScenarioParsingError, + FileNotFoundError, + OSError, + ValueError, + ) as exc: + logging.error(str(exc)) + return 1 + + register_signal_handlers(runner.cancel_on_signal) + try: + successful = cloudai.handlers.execute_experiment( + runner, tests, enable_cache_without_check=args.enable_cache_without_check + ) + except cloudai.handlers.InstallationError: + return 1 + + if args.single_sbatch or not any(tr.is_dse_job for tr in scenario.test_runs): + return 0 + return 0 if successful else 1 + + def setup_logging(log_file: str, log_level: str) -> None: """ Configure logging for the application. @@ -195,7 +242,7 @@ def dry_run( enable_cache_without_check=enable_cache_without_check, single_sbatch=single_sbatch, ) - exit(handle_dry_run_and_run(args)) + exit(_handle_dry_run_and_run(args)) @main.command() @@ -227,7 +274,7 @@ def run( enable_cache_without_check=enable_cache_without_check, single_sbatch=single_sbatch, ) - exit(handle_dry_run_and_run(args)) + exit(_handle_dry_run_and_run(args)) @main.command() diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 09402d619..d7c5f3ead 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -17,8 +17,9 @@ import argparse import copy import logging -import signal +import tempfile import traceback +import warnings from contextlib import contextmanager from pathlib import Path from typing import Callable, List, Optional @@ -32,8 +33,6 @@ BaseInstaller, CloudAIGymEnv, Installable, - InstallStatusResult, - MissingTestError, Parser, Registry, Runner, @@ -131,6 +130,10 @@ def _scenario_installables(scenario: TestScenario) -> list[Installable]: def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: + return _run_dse_job(runner, args.mode) + + +def _run_dse_job(runner: Runner, mode: str) -> int: registry = Registry() original_test_runs = copy.deepcopy(runner.runner.test_scenario.test_runs) @@ -182,7 +185,7 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: run_error = exc logging.exception("DSE job aborted by an unexpected error; generating reports before failing.") - if args.mode == "run": + if mode == "run": runner.runner.test_scenario.test_runs = original_test_runs generate_reports( runner.runner.system, @@ -243,7 +246,7 @@ def generate_reports( logging.debug(e, exc_info=True) -def handle_non_dse_job(runner: Runner, args: argparse.Namespace) -> bool: +def handle_non_dse_job(runner: Runner, args: argparse.Namespace | None = None) -> bool: successful = runner.run() generate_reports(runner.runner.system, runner.runner.test_scenario, runner.runner.scenario_root) logging.info("All jobs are complete.") @@ -251,129 +254,116 @@ def handle_non_dse_job(runner: Runner, args: argparse.Namespace) -> bool: 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 + """Compatibility entry point for CLI signal registration.""" + warnings.warn("Use cloudai.cli.cli.register_signal_handlers instead.", DeprecationWarning, stacklevel=2) + from cloudai.cli.cli import register_signal_handlers as cli_register_signal_handlers + cli_register_signal_handlers(signal_handler) -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 +class InstallationError(RuntimeError): + """Workload prerequisites could not be installed.""" - Registry().update_runner("slurm", SingleSbatchRunner) - return True +def load_experiment( + scenario_path: Path, + system_path: Path, + *, + tests_dir: Path | None = None, + hook_dir: Path | None = None, +) -> tuple[System, list[TestDefinition], TestScenario]: + """Parse a scenario for programmatic execution without process exits.""" + parser = Parser(system_path, hook_dir or HOOK_ROOT, exit_on_error=False) + system, tests, scenario = parser.parse(tests_dir, scenario_path) + if scenario is None: + raise TestScenarioParsingError("A test scenario is required.") + validate_domain_randomization_active(scenario) + return system, tests, scenario -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 +def create_experiment_runner( + system: System, + scenario: TestScenario, + *, + mode: str = "run", + output_dir: Path | None = None, + result_dir: Path | None = None, + single_sbatch: bool = False, + unique_result_dir: bool = False, +) -> Runner: + """Prepare the system and create a runner for one experiment.""" + if single_sbatch and not isinstance(system, SlurmSystem): + raise ValueError("Single sbatch is only supported for Slurm systems.") + if output_dir is not None: + system.output_path = output_dir.absolute() + if not prepare_output_dir(system.output_path): + raise OSError(f"Cannot prepare output directory: {system.output_path}") + if mode == "dry-run": + system.monitor_interval = 1 + system.update() + if unique_result_dir: + result_dir = Path(tempfile.mkdtemp(prefix=f"{scenario.name}_", dir=system.output_path)).resolve() + return Runner( + mode, + system, + scenario, + runner_class=SingleSbatchRunner if single_sbatch else None, + output_path=result_dir, + ) + + +def execute_experiment( + runner: Runner, tests: list[TestDefinition], *, enable_cache_without_check: bool = False +) -> bool: + """Install prerequisites, execute the scenario, and finalize its output.""" + system = runner.runner.system + scenario = runner.runner.test_scenario + mode = runner.runner.mode + successful = False 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: + runner.runner.experiment_output.write() + logging.info(f"System Name: {system.name}") + logging.info(f"Scheduler: {system.scheduler}") + logging.info(f"Test Scenario Name: {scenario.name}") + logging.info("Checking if workloads components are installed.") + installables, installer = prepare_installation(system, tests, scenario) + if enable_cache_without_check: + result = installer.mark_as_installed(installables) + else: + result = installer.is_installed(installables) + if mode == "run" and not result.success: + logging.info("Not all workloads components are installed. Installing...") + result = installer.install(installables) + if not result.success: + logging.error("Failed to install workloads components.") + logging.error(result.message) + raise InstallationError(result.message) _log_installation_dirs("CloudAI is successfully installed into", system) + elif mode == "dry-run": + result = installer.mark_as_installed(installables) + if not result.success: + logging.warning("Failed to mark workloads components as installed for dry-run.") + + logging.info(scenario.pretty_print()) + logging.info(f"Scenario results will be stored at: {runner.runner.scenario_root}") + has_dse = any(tr.is_dse_job for tr in scenario.test_runs) + if isinstance(runner.runner, SingleSbatchRunner) or not has_dse: + successful = handle_non_dse_job(runner) + elif all(tr.is_dse_job for tr in scenario.test_runs): + successful = _run_dse_job(runner, mode) == 0 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}") + logging.error("Mixing DSE and non-DSE jobs is not allowed.") + return successful + finally: + runner.runner.finish_output(successful) - 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 +def handle_dry_run_and_run(args: argparse.Namespace) -> int: + """Compatibility entry point for the CLI run commands.""" + warnings.warn("Use cloudai.api.run_experiment instead.", DeprecationWarning, stacklevel=2) + from cloudai.cli.cli import _handle_dry_run_and_run - logging.error("Mixing DSE and non-DSE jobs is not allowed.") - return 1 - finally: - runner.runner.finish_output(successful) + return _handle_dry_run_and_run(args) def handle_generate_report(args: argparse.Namespace) -> int: diff --git a/src/cloudai/parser.py b/src/cloudai/parser.py index b76dc6787..3fe02ddeb 100644 --- a/src/cloudai/parser.py +++ b/src/cloudai/parser.py @@ -42,16 +42,18 @@ class Parser: """Main parser for parsing all types of configurations.""" - def __init__(self, system_config_path: Path, hook_root: Path = HOOK_ROOT) -> None: + def __init__(self, system_config_path: Path, hook_root: Path = HOOK_ROOT, *, exit_on_error: bool = True) -> None: """ Initialize a Parser instance. Args: system_config_path (str): The file path for system configurations. hook_root (Path): Directory containing hook scenarios and tests. + exit_on_error (bool): Exit on parsing errors instead of raising them. """ logging.debug(f"Initializing parser with: {system_config_path=}") self.system_config_path = system_config_path + self.exit_on_error = exit_on_error self.hook_root = hook_root self.hook_test_root = hook_root / "test" self._system: Optional[System] = None @@ -64,6 +66,8 @@ def system(self) -> System: try: self._system = self.parse_system(self.system_config_path) except SystemConfigParsingError: + if not self.exit_on_error: + raise exit(1) # exit right away to keep error message readable for users return self._system @@ -82,27 +86,31 @@ def parse( Tuple[System, List[Test], Optional[TestScenario]]: A tuple containing the system object, a list of filtered test template objects, and the main test scenario object if provided. """ + try: + return self._parse(test_path, test_scenario_path) + except (SystemConfigParsingError, TestConfigParsingError, TestScenarioParsingError): + if not self.exit_on_error: + raise + raise SystemExit(1) from None + + def _parse( + self, test_path: Path | None, test_scenario_path: Path | None + ) -> tuple[System, list[TestDefinition], TestScenario | None]: tests: list[TestDefinition] = [] if test_path: if not test_path.exists(): raise FileNotFoundError(f"Test path '{test_path}' not found.") - try: - tests = self.parse_tests(list(test_path.glob("*.toml")), self.system) - except TestConfigParsingError: - exit(1) # exit right away to keep error message readable for users + tests = self.parse_tests(list(test_path.glob("*.toml")), self.system) if not self.hook_root.exists(): logging.debug(f"Hook root path '{self.hook_root}' does not exist.") - try: - hook_tests = ( - self.parse_tests(list(self.hook_test_root.glob("*.toml")), self.system) - if self.hook_test_root.exists() - else [] - ) - except TestConfigParsingError: - exit(1) # exit right away to keep error message readable for users + hook_tests = ( + self.parse_tests(list(self.hook_test_root.glob("*.toml")), self.system) + if self.hook_test_root.exists() + else [] + ) if not test_scenario_path: all_tests = list({test.name: test for test in tests + hook_tests}.values()) @@ -111,19 +119,13 @@ def parse( test_mapping = {t.name: t for t in tests} hook_test_scenario_mapping = {} if self.hook_root.exists() and list(self.hook_root.glob("*.toml")): - try: - hook_test_scenario_mapping = self.parse_hooks( - list(self.hook_root.glob("*.toml")), self.system, {t.name: t for t in hook_tests} - ) - except TestScenarioParsingError: - exit(1) # exit right away to keep error message readable for users - - try: - test_scenario = self.parse_test_scenario( - test_scenario_path, self.system, test_mapping, hook_test_scenario_mapping + hook_test_scenario_mapping = self.parse_hooks( + list(self.hook_root.glob("*.toml")), self.system, {t.name: t for t in hook_tests} ) - except TestScenarioParsingError: - exit(1) # exit right away to keep error message readable for users + + test_scenario = self.parse_test_scenario( + test_scenario_path, self.system, test_mapping, hook_test_scenario_mapping + ) test_scenario.config_paths = ConfigPaths( system_path=self.system_config_path.resolve(), diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index aa65b2537..0ba9dda3d 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -25,10 +25,10 @@ import pytest import toml +import cloudai.cli.cli import cloudai.models.output from cloudai.cli import setup_logging 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 ( @@ -149,7 +149,7 @@ def do_dry_run(self, tmp_path_factory: pytest.TempPathFactory, request: pytest.F mock_process.communicate.return_value = ("", "") mock_execute.return_value = mock_process - handle_dry_run_and_run(args) + cloudai.cli.cli._handle_dry_run_and_run(args) return (tmp_path, scenario) diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 000000000..9c76d61ff --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,87 @@ +# 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 +import signal + +import pytest +import toml + +import cloudai.api +import cloudai.core +from cloudai.systems.standalone import StandaloneSystem +from cloudai.systems.standalone.standalone_job import StandaloneJob +from cloudai.systems.standalone.standalone_runner import StandaloneRunner +from cloudai.workloads.sleep import SleepTestDefinition + + +def test_run_and_list_experiments( + standalone_system: StandaloneSystem, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch +) -> None: + system = tmp_path / "system.toml" + system.write_text( + toml.dumps( + { + "name": standalone_system.name, + "scheduler": "standalone", + "install_path": str(standalone_system.install_path), + "output_path": str(standalone_system.output_path), + "monitor_interval": 0, + } + ) + ) + scenario = tmp_path / "scenario.toml" + scenario.write_text( + toml.dumps( + { + "name": "api-test", + "Tests": [ + { + "id": "sleep", + "name": "sleep", + "description": "API smoke test", + "test_template_name": "Sleep", + "cmd_args": {"seconds": 0}, + } + ], + } + ) + ) + monkeypatch.setattr(StandaloneSystem, "is_job_completed", lambda *_: True) + monkeypatch.setattr(StandaloneSystem, "is_job_running", lambda *_: False) + monkeypatch.setattr( + StandaloneRunner, + "_submit_test", + lambda self, tr: StandaloneJob(tr, id=123, start=datetime.datetime.now(datetime.timezone.utc)), + ) + monkeypatch.setattr( + SleepTestDefinition, + "was_run_successful", + lambda *_: cloudai.core.JobStatusResult(True, "test outcome"), + ) + + signal_handler = signal.getsignal(signal.SIGINT) + assert cloudai.api.list_experiments(system) == [] + experiment = cloudai.api.run_experiment(scenario, system) + assert experiment.status == "completed" + assert experiment.tests[0].runs[0].status == "completed" + assert cloudai.api.list_experiments(system) == [(experiment.id, pathlib.Path(experiment.path))] + assert signal.getsignal(signal.SIGINT) == signal_handler + + scenario.write_text('name = "invalid"\n') + with pytest.raises(cloudai.core.TestScenarioParsingError): + cloudai.api.run_experiment(scenario, system) From eb304dc6ef22fc848803cbe6ee59ff9293264bdc Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Wed, 30 Sep 2026 14:06:35 +0200 Subject: [PATCH 02/13] Keep legacy run signatures in CLI compatibility shims Signed-off-by: Ivan Podkidyshev --- src/cloudai/cli/cli.py | 6 +++--- src/cloudai/cli/handlers.py | 36 ++++++++++++++++++++++++++++++------ src/cloudai/handlers.py | 31 +++++-------------------------- tests/test_acceptance.py | 2 +- tests/test_handlers.py | 19 +++++++++---------- 5 files changed, 48 insertions(+), 46 deletions(-) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 93b5488f6..d68f5c958 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -39,7 +39,7 @@ def register_signal_handlers(signal_handler: Callable) -> None: signal.signal(sig, signal_handler) -def _handle_dry_run_and_run(args: argparse.Namespace) -> int: +def handle_dry_run_and_run(args: argparse.Namespace) -> int: """Translate CLI options and experiment outcomes into a process exit code.""" try: system, tests, scenario = cloudai.handlers.load_experiment( @@ -242,7 +242,7 @@ def dry_run( enable_cache_without_check=enable_cache_without_check, single_sbatch=single_sbatch, ) - exit(_handle_dry_run_and_run(args)) + exit(handle_dry_run_and_run(args)) @main.command() @@ -274,7 +274,7 @@ def run( enable_cache_without_check=enable_cache_without_check, single_sbatch=single_sbatch, ) - exit(_handle_dry_run_and_run(args)) + exit(handle_dry_run_and_run(args)) @main.command() diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index de41076f9..29f443394 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -14,17 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import argparse import functools +import types import typing import warnings +import cloudai.cli.cli +import cloudai.core import cloudai.handlers _P = typing.ParamSpec("_P") _R = typing.TypeVar("_R") -def _deprecated(function: typing.Callable[_P, _R]) -> typing.Callable[_P, _R]: +def _deprecated( + function: typing.Callable[_P, _R], owner: types.ModuleType = cloudai.handlers +) -> typing.Callable[_P, _R]: @functools.wraps(function) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: warnings.warn( @@ -32,18 +38,16 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: DeprecationWarning, stacklevel=2, ) - return getattr(cloudai.handlers, function.__name__)(*args, **kwargs) + return getattr(owner, function.__name__)(*args, **kwargs) return wrapper handle_install_and_uninstall = _deprecated(cloudai.handlers.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) +register_signal_handlers = _deprecated(cloudai.cli.cli.register_signal_handlers, cloudai.cli.cli) +handle_dry_run_and_run = _deprecated(cloudai.cli.cli.handle_dry_run_and_run, cloudai.cli.cli) 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) @@ -56,3 +60,23 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: 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) + + +def handle_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> int: + """Retain the former CLI call shape while using the runner's mode.""" + warnings.warn( + "cloudai.cli.handlers.handle_dse_job is deprecated; do not call CLI handlers directly.", + DeprecationWarning, + stacklevel=2, + ) + return cloudai.handlers.handle_dse_job(runner) + + +def handle_non_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> bool: + """Retain the former CLI call shape without passing unused arguments onward.""" + warnings.warn( + "cloudai.cli.handlers.handle_non_dse_job is deprecated; do not call CLI handlers directly.", + DeprecationWarning, + stacklevel=2, + ) + return cloudai.handlers.handle_non_dse_job(runner) diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index d7c5f3ead..cd75c1116 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -19,10 +19,9 @@ import logging import tempfile import traceback -import warnings from contextlib import contextmanager from pathlib import Path -from typing import Callable, List, Optional +from typing import List, Optional from unittest.mock import Mock import toml @@ -129,11 +128,7 @@ def _scenario_installables(scenario: TestScenario) -> list[Installable]: return installables -def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: - return _run_dse_job(runner, args.mode) - - -def _run_dse_job(runner: Runner, mode: str) -> int: +def handle_dse_job(runner: Runner) -> int: registry = Registry() original_test_runs = copy.deepcopy(runner.runner.test_scenario.test_runs) @@ -185,7 +180,7 @@ def _run_dse_job(runner: Runner, mode: str) -> int: run_error = exc logging.exception("DSE job aborted by an unexpected error; generating reports before failing.") - if mode == "run": + if runner.runner.mode == "run": runner.runner.test_scenario.test_runs = original_test_runs generate_reports( runner.runner.system, @@ -246,21 +241,13 @@ def generate_reports( logging.debug(e, exc_info=True) -def handle_non_dse_job(runner: Runner, args: argparse.Namespace | None = None) -> bool: +def handle_non_dse_job(runner: Runner) -> bool: successful = runner.run() generate_reports(runner.runner.system, runner.runner.test_scenario, runner.runner.scenario_root) logging.info("All jobs are complete.") return successful -def register_signal_handlers(signal_handler: Callable) -> None: - """Compatibility entry point for CLI signal registration.""" - warnings.warn("Use cloudai.cli.cli.register_signal_handlers instead.", DeprecationWarning, stacklevel=2) - from cloudai.cli.cli import register_signal_handlers as cli_register_signal_handlers - - cli_register_signal_handlers(signal_handler) - - class InstallationError(RuntimeError): """Workload prerequisites could not be installed.""" @@ -350,7 +337,7 @@ def execute_experiment( if isinstance(runner.runner, SingleSbatchRunner) or not has_dse: successful = handle_non_dse_job(runner) elif all(tr.is_dse_job for tr in scenario.test_runs): - successful = _run_dse_job(runner, mode) == 0 + successful = handle_dse_job(runner) == 0 else: logging.error("Mixing DSE and non-DSE jobs is not allowed.") return successful @@ -358,14 +345,6 @@ def execute_experiment( runner.runner.finish_output(successful) -def handle_dry_run_and_run(args: argparse.Namespace) -> int: - """Compatibility entry point for the CLI run commands.""" - warnings.warn("Use cloudai.api.run_experiment instead.", DeprecationWarning, stacklevel=2) - from cloudai.cli.cli import _handle_dry_run_and_run - - return _handle_dry_run_and_run(args) - - def handle_generate_report(args: argparse.Namespace) -> int: """ Generate a report based on the existing configuration and test results. diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 0ba9dda3d..3857169cd 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -149,7 +149,7 @@ def do_dry_run(self, tmp_path_factory: pytest.TempPathFactory, request: pytest.F mock_process.communicate.return_value = ("", "") mock_execute.return_value = mock_process - cloudai.cli.cli._handle_dry_run_and_run(args) + cloudai.cli.cli.handle_dry_run_and_run(args) return (tmp_path, scenario) diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 1e78d117f..5fd76a049 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse import copy from pathlib import Path from typing import Any, ClassVar, Iterator, Optional @@ -114,7 +113,7 @@ def test_dse_run_does_not_support_dependencies( dse_tr.dependencies = {dep: TestDependency(test_run=dse_tr)} test_scenario: TestScenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 1 + assert handle_dse_job(runner) == 1 assert "Dependencies are not supported for DSE jobs, all cases run consecutively." in caplog.text assert "Please remove dependencies and re-run." in caplog.text @@ -160,7 +159,7 @@ def test_dse_run_uses_agent_config( test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 0 + assert handle_dse_job(runner) == 0 assert len(StubAgent.received_configs) == 1 recorded = StubAgent.received_configs[0] @@ -255,7 +254,7 @@ def _job_output_path(tr: TestRun, create: bool = True): # run test with caplog.at_level("INFO"): - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 0 + assert handle_dse_job(runner) == 0 reporter = StatusReporter( inner_runner.system, @@ -397,7 +396,7 @@ def test_handle_dse_job_invokes_agent_run( update_output = MagicMock() monkeypatch.setattr(CloudAIGymEnv, "update_output", update_output) - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 0 + assert handle_dse_job(runner) == 0 assert CustomRunStubAgent.run_calls == 1 assert update_output.call_count == 2 @@ -413,7 +412,7 @@ def test_handle_dse_job_propagates_agent_run_nonzero_rc( test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 1 + assert handle_dse_job(runner) == 1 assert CustomRunStubAgent.run_calls == 1 @@ -435,7 +434,7 @@ def test_handle_dse_job_accumulates_nonzero_rc_and_continues( test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr, second_tr]) runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 1 + assert handle_dse_job(runner) == 1 assert CustomRunStubAgent.run_calls == 2 @@ -458,7 +457,7 @@ def test_handle_dse_job_propagates_agent_run_exception( monkeypatch.setattr(CloudAIGymEnv, "update_output", update_output) with pytest.raises(RuntimeError, match="agent blew up"): - handle_dse_job(runner, argparse.Namespace(mode="dry-run")) + handle_dse_job(runner) assert CustomRunStubAgent.run_calls == 1 assert update_output.call_count == 2 @@ -477,7 +476,7 @@ def test_handle_dse_job_hard_fail_aborts_remaining_runs( runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) with pytest.raises(RuntimeError, match="agent blew up"): - handle_dse_job(runner, argparse.Namespace(mode="dry-run")) + handle_dse_job(runner) assert CustomRunStubAgent.run_calls == 1 @@ -495,7 +494,7 @@ def test_handle_dse_job_documents_failure_in_reports_before_raising( runner.runner.scenario_root = tmp_path with pytest.raises(RuntimeError, match="agent blew up"): - handle_dse_job(runner, argparse.Namespace(mode="run")) + handle_dse_job(runner) failure_report = tmp_path / "dse_failure.txt" assert failure_report.exists() From 749d08ae4e455a1d1488f5824ffaaeca830d93ca Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Wed, 30 Sep 2026 14:30:30 +0200 Subject: [PATCH 03/13] Add experiment lookup to Python API Signed-off-by: Ivan Podkidyshev --- doc/api.rst | 5 +++++ src/cloudai/api.py | 15 +++++++++++++++ tests/test_api.py | 7 +++++++ 3 files changed, 27 insertions(+) diff --git a/doc/api.rst b/doc/api.rst index 8f6c5c3de..389a71361 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -27,3 +27,8 @@ use a custom hook directory. from the system's configured output directory. Directories without ``experiment.json`` are skipped. A missing output directory gives an empty list; an invalid experiment file raises an error. + +``get_experiment(experiment_id, system)`` loads an ``Experiment`` from that +system's output directory. Pass a ``Path`` to a result directory or its +``experiment.json`` file to load it directly. A missing file raises +``FileNotFoundError``. diff --git a/src/cloudai/api.py b/src/cloudai/api.py index b5dca98ce..e8db697c6 100644 --- a/src/cloudai/api.py +++ b/src/cloudai/api.py @@ -58,3 +58,18 @@ def list_experiments(system: pathlib.Path) -> list[tuple[str, pathlib.Path]]: experiment = cloudai.models.output.Experiment.model_validate_json(path.read_text(encoding="utf-8")) experiments.append((experiment.id, path.parent)) return experiments + + +def get_experiment(exp_id: str | pathlib.Path, system: str | pathlib.Path) -> cloudai.models.output.Experiment: + """Load an experiment by ID or by its result directory or JSON path.""" + if isinstance(exp_id, pathlib.Path): + path = exp_id.expanduser() + if path.suffix != ".json": + path /= "experiment.json" + else: + if exp_id in ("", ".", "..") or pathlib.Path(exp_id).name != exp_id: + raise ValueError(f"Invalid experiment ID: {exp_id!r}") + output_dir = Parser.parse_system(pathlib.Path(system).expanduser().resolve()).output_path.expanduser().resolve() + path = output_dir / exp_id / "experiment.json" + + return cloudai.models.output.Experiment.model_validate_json(path.read_text(encoding="utf-8")) diff --git a/tests/test_api.py b/tests/test_api.py index 9c76d61ff..5625b92a7 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -80,6 +80,13 @@ def test_run_and_list_experiments( assert experiment.status == "completed" assert experiment.tests[0].runs[0].status == "completed" assert cloudai.api.list_experiments(system) == [(experiment.id, pathlib.Path(experiment.path))] + assert cloudai.api.get_experiment(experiment.id, str(system)) == experiment + assert cloudai.api.get_experiment(pathlib.Path(experiment.path), system) == experiment + assert cloudai.api.get_experiment(pathlib.Path(experiment.path) / "experiment.json", system) == experiment + with pytest.raises(FileNotFoundError): + cloudai.api.get_experiment("missing", system) + with pytest.raises(ValueError, match="Invalid experiment ID"): + cloudai.api.get_experiment("../other", system) assert signal.getsignal(signal.SIGINT) == signal_handler scenario.write_text('name = "invalid"\n') From c129fad226f94007528ef16fcc878d59d46f2e6b Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 15:25:37 +0200 Subject: [PATCH 04/13] Simplify API experiment discovery Signed-off-by: Ivan Podkidyshev --- doc/api.rst | 2 +- src/cloudai/_core/runner.py | 19 ++++++++++--------- src/cloudai/api.py | 12 ++++++------ src/cloudai/cli/handlers.py | 9 ++++++--- src/cloudai/handlers.py | 14 ++------------ tests/test_api.py | 10 +++++++++- tests/test_cloudaigym.py | 4 ++-- tests/test_handlers.py | 24 +++++++++++++++--------- 8 files changed, 51 insertions(+), 43 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index 389a71361..5b6ea3999 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -26,7 +26,7 @@ use a custom hook directory. ``list_experiments(system)`` returns ``(experiment_id, result_directory)`` pairs from the system's configured output directory. Directories without ``experiment.json`` are skipped. A missing output directory gives an empty list; -an invalid experiment file raises an error. +experiment file contents are not parsed. ``get_experiment(experiment_id, system)`` loads an ``Experiment`` from that system's output directory. Pass a ``Path`` to a result directory or its diff --git a/src/cloudai/_core/runner.py b/src/cloudai/_core/runner.py index 17552c6bd..d2b150ae0 100644 --- a/src/cloudai/_core/runner.py +++ b/src/cloudai/_core/runner.py @@ -16,7 +16,7 @@ import datetime import logging -from pathlib import Path +import warnings from types import FrameType from typing import Optional @@ -52,15 +52,18 @@ def __init__( test_scenario: TestScenario, *, runner_class: type[BaseRunner] | None = None, - output_path: Path | None = None, ): logging.info(f"Initializing Runner [{mode.upper()}] mode") - if runner_class is None and output_path is None: + if runner_class is None: + warnings.warn( + "Runner.create_runner() compatibility with its former three-argument signature is deprecated and " + "will be removed in CloudAI 1.9.1.", + DeprecationWarning, + stacklevel=2, + ) self.runner = self.create_runner(mode, system, test_scenario) else: - self.runner = self.create_runner( - mode, system, test_scenario, runner_class=runner_class, output_path=output_path - ) + self.runner = self.create_runner(mode, system, test_scenario, runner_class=runner_class) def create_runner( self, @@ -69,7 +72,6 @@ def create_runner( test_scenario: TestScenario, *, runner_class: type[BaseRunner] | None = None, - output_path: Path | None = None, ) -> BaseRunner: """ Dynamically create a runner instance based on the system's scheduler type. @@ -79,7 +81,6 @@ def create_runner( system (System): The system configuration. test_scenario (TestScenario): The test scenario to run. runner_class (type[BaseRunner] | None): Override for this invocation. - output_path (Path | None): Exact scenario result directory, if allocated. Returns: BaseRunner: A runner instance suitable for the system. @@ -99,7 +100,7 @@ def create_runner( if not system.output_path.exists(): system.output_path.mkdir() current_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - results_root = output_path or system.output_path / f"{test_scenario.name}_{current_time}" + results_root = system.output_path / f"{test_scenario.name}_{current_time}" return runner_class(mode, system, test_scenario, results_root) diff --git a/src/cloudai/api.py b/src/cloudai/api.py index e8db697c6..88874f570 100644 --- a/src/cloudai/api.py +++ b/src/cloudai/api.py @@ -40,7 +40,7 @@ def run_experiment( tests_dir=tests_dir, hook_dir=hook_dir, ) - runner = cloudai.handlers.create_experiment_runner(parsed_system, parsed_scenario, unique_result_dir=True) + runner = cloudai.handlers.create_experiment_runner(parsed_system, parsed_scenario) cloudai.handlers.execute_experiment(runner, tests) return runner.runner.experiment_output.snapshot() @@ -53,11 +53,11 @@ def list_experiments(system: pathlib.Path) -> list[tuple[str, pathlib.Path]]: if not output_dir.is_dir(): raise NotADirectoryError(output_dir) - experiments = [] - for path in sorted(output_dir.glob("*/experiment.json")): - experiment = cloudai.models.output.Experiment.model_validate_json(path.read_text(encoding="utf-8")) - experiments.append((experiment.id, path.parent)) - return experiments + return [ + (experiment_file.parent.name, experiment_file.parent) + for experiment_file in sorted(output_dir.glob("*/experiment.json")) + if experiment_file.is_file() + ] def get_experiment(exp_id: str | pathlib.Path, system: str | pathlib.Path) -> cloudai.models.output.Experiment: diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 29f443394..6d5b105df 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -34,7 +34,8 @@ def _deprecated( @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.", + f"cloudai.cli.handlers.{function.__name__} is deprecated; do not call CLI handlers directly. " + "This compatibility shim will be removed in CloudAI 1.9.1.", DeprecationWarning, stacklevel=2, ) @@ -65,7 +66,8 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: def handle_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> int: """Retain the former CLI call shape while using the runner's mode.""" warnings.warn( - "cloudai.cli.handlers.handle_dse_job is deprecated; do not call CLI handlers directly.", + "cloudai.cli.handlers.handle_dse_job is deprecated; do not call CLI handlers directly. " + "This compatibility shim will be removed in CloudAI 1.9.1.", DeprecationWarning, stacklevel=2, ) @@ -75,7 +77,8 @@ def handle_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> int def handle_non_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> bool: """Retain the former CLI call shape without passing unused arguments onward.""" warnings.warn( - "cloudai.cli.handlers.handle_non_dse_job is deprecated; do not call CLI handlers directly.", + "cloudai.cli.handlers.handle_non_dse_job is deprecated; do not call CLI handlers directly. " + "This compatibility shim will be removed in CloudAI 1.9.1.", DeprecationWarning, stacklevel=2, ) diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index cd75c1116..1575e68c8 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -17,7 +17,6 @@ import argparse import copy import logging -import tempfile import traceback from contextlib import contextmanager from pathlib import Path @@ -274,9 +273,7 @@ def create_experiment_runner( *, mode: str = "run", output_dir: Path | None = None, - result_dir: Path | None = None, single_sbatch: bool = False, - unique_result_dir: bool = False, ) -> Runner: """Prepare the system and create a runner for one experiment.""" if single_sbatch and not isinstance(system, SlurmSystem): @@ -288,15 +285,8 @@ def create_experiment_runner( if mode == "dry-run": system.monitor_interval = 1 system.update() - if unique_result_dir: - result_dir = Path(tempfile.mkdtemp(prefix=f"{scenario.name}_", dir=system.output_path)).resolve() - return Runner( - mode, - system, - scenario, - runner_class=SingleSbatchRunner if single_sbatch else None, - output_path=result_dir, - ) + runner_class = SingleSbatchRunner if single_sbatch else Registry().runners_map.get(system.scheduler.lower()) + return Runner(mode, system, scenario, runner_class=runner_class) def execute_experiment( diff --git a/tests/test_api.py b/tests/test_api.py index 5625b92a7..b1fd0218d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -79,7 +79,15 @@ def test_run_and_list_experiments( experiment = cloudai.api.run_experiment(scenario, system) assert experiment.status == "completed" assert experiment.tests[0].runs[0].status == "completed" - assert cloudai.api.list_experiments(system) == [(experiment.id, pathlib.Path(experiment.path))] + experiment_path = pathlib.Path(experiment.path) + assert experiment.id == experiment_path.name + unparsed_path = standalone_system.output_path / "zzz-unparsed" + unparsed_path.mkdir() + (unparsed_path / "experiment.json").write_text("not JSON") + assert cloudai.api.list_experiments(system) == [ + (experiment.id, experiment_path), + (unparsed_path.name, unparsed_path.resolve()), + ] assert cloudai.api.get_experiment(experiment.id, str(system)) == experiment assert cloudai.api.get_experiment(pathlib.Path(experiment.path), system) == experiment assert cloudai.api.get_experiment(pathlib.Path(experiment.path) / "experiment.json", system) == experiment diff --git a/tests/test_cloudaigym.py b/tests/test_cloudaigym.py index 4dc5cd1b3..a7aa60720 100644 --- a/tests/test_cloudaigym.py +++ b/tests/test_cloudaigym.py @@ -28,7 +28,7 @@ ) from cloudai.configurator.env_params import EnvParamSpec, ObsLeafDescriptor from cloudai.core import BaseRunner, RewardOverrides, Runner, TestRun, TestScenario -from cloudai.systems.slurm import SlurmSystem +from cloudai.systems.slurm import SlurmRunner, SlurmSystem from cloudai.util import flatten_dict from cloudai.workloads.nemo_run import ( Data, @@ -97,7 +97,7 @@ def setup_env(slurm_system: SlurmSystem, nemorun: NeMoRunTestDefinition) -> tupl slurm_system.output_path / test_scenario.name / test_run.name / f"{test_run.current_iteration}" ) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) return test_run, runner.runner diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 5fd76a049..71ae9fff1 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -49,7 +49,7 @@ from cloudai.models.scenario import ReportConfig from cloudai.models.workload import CmdArgs, TestDefinition from cloudai.reporter import StatusReporter -from cloudai.systems.slurm.slurm_system import SlurmSystem +from cloudai.systems.slurm import SlurmRunner, SlurmSystem class StubAgentConfig(BaseAgentConfig): @@ -96,6 +96,12 @@ def stub_agent_name() -> Iterator[str]: registry.update_agent(agent_name, old_agent) +def test_runner_warns_about_legacy_create_runner_signature(slurm_system: SlurmSystem, dse_tr: TestRun) -> None: + scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) + with pytest.warns(DeprecationWarning, match="will be removed in CloudAI 1.9.1"): + Runner(mode="dry-run", system=slurm_system, test_scenario=scenario) + + @pytest.mark.parametrize("dep", ["start_post_comp", "start_post_init", "end_post_comp"]) def test_dse_run_does_not_support_dependencies( slurm_system: SlurmSystem, dse_tr: TestRun, dep: str, caplog: pytest.LogCaptureFixture @@ -112,7 +118,7 @@ def test_dse_run_does_not_support_dependencies( """ dse_tr.dependencies = {dep: TestDependency(test_run=dse_tr)} test_scenario: TestScenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) assert handle_dse_job(runner) == 1 assert "Dependencies are not supported for DSE jobs, all cases run consecutively." in caplog.text assert "Please remove dependencies and re-run." in caplog.text @@ -157,7 +163,7 @@ def test_dse_run_uses_agent_config( dse_tr.test.agent = stub_agent_name dse_tr.test.agent_config = agent_config test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) assert handle_dse_job(runner) == 0 assert len(StubAgent.received_configs) == 1 @@ -392,7 +398,7 @@ def test_handle_dse_job_invokes_agent_run( """``handle_dse_job`` must delegate orchestration to ``agent.run()`` (polymorphism).""" dse_tr.test.agent = custom_run_agent_name test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) update_output = MagicMock() monkeypatch.setattr(CloudAIGymEnv, "update_output", update_output) @@ -410,7 +416,7 @@ def test_handle_dse_job_propagates_agent_run_nonzero_rc( CustomRunStubAgent.run_returns = 1 dse_tr.test.agent = custom_run_agent_name test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) assert handle_dse_job(runner) == 1 assert CustomRunStubAgent.run_calls == 1 @@ -432,7 +438,7 @@ def test_handle_dse_job_accumulates_nonzero_rc_and_continues( second_tr = copy.deepcopy(dse_tr) second_tr.name = "dse_second" test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr, second_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) assert handle_dse_job(runner) == 1 assert CustomRunStubAgent.run_calls == 2 @@ -452,7 +458,7 @@ def test_handle_dse_job_propagates_agent_run_exception( CustomRunStubAgent.run_raises = RuntimeError("agent blew up") dse_tr.test.agent = custom_run_agent_name test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) update_output = MagicMock() monkeypatch.setattr(CloudAIGymEnv, "update_output", update_output) @@ -473,7 +479,7 @@ def test_handle_dse_job_hard_fail_aborts_remaining_runs( second_tr = copy.deepcopy(dse_tr) second_tr.name = "dse_second" test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr, second_tr]) - runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) with pytest.raises(RuntimeError, match="agent blew up"): handle_dse_job(runner) @@ -490,7 +496,7 @@ def test_handle_dse_job_documents_failure_in_reports_before_raising( CustomRunStubAgent.run_raises = RuntimeError("agent blew up") dse_tr.test.agent = custom_run_agent_name test_scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) - runner = Runner(mode="run", system=slurm_system, test_scenario=test_scenario) + runner = Runner(mode="run", system=slurm_system, test_scenario=test_scenario, runner_class=SlurmRunner) runner.runner.scenario_root = tmp_path with pytest.raises(RuntimeError, match="agent blew up"): From b0a861799b7771943deec4460a7c7cd2474c17ff Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 15:43:58 +0200 Subject: [PATCH 05/13] Constrain experiment lookup to result directories Signed-off-by: Ivan Podkidyshev --- doc/api.rst | 6 +++--- src/cloudai/api.py | 37 ++++++++++++++++++++++++------------- tests/test_api.py | 7 +++++-- 3 files changed, 32 insertions(+), 18 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index 5b6ea3999..138324cd1 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -29,6 +29,6 @@ from the system's configured output directory. Directories without experiment file contents are not parsed. ``get_experiment(experiment_id, system)`` loads an ``Experiment`` from that -system's output directory. Pass a ``Path`` to a result directory or its -``experiment.json`` file to load it directly. A missing file raises -``FileNotFoundError``. +system's output directory. Pass a ``Path`` to a result directory to load it +directly. Invalid IDs and missing, unreadable, or malformed experiment files +raise ``ValueError``. diff --git a/src/cloudai/api.py b/src/cloudai/api.py index 88874f570..de04e59cc 100644 --- a/src/cloudai/api.py +++ b/src/cloudai/api.py @@ -16,9 +16,11 @@ import pathlib +import pydantic + +import cloudai.core import cloudai.handlers import cloudai.models.output -from cloudai.core import Parser def run_experiment( @@ -47,7 +49,7 @@ def run_experiment( def list_experiments(system: pathlib.Path) -> list[tuple[str, pathlib.Path]]: """List saved experiment IDs and result directories for a system.""" - output_dir = Parser.parse_system(system.expanduser().resolve()).output_path.expanduser().resolve() + output_dir = cloudai.core.Parser.parse_system(system.expanduser().resolve()).output_path.expanduser().resolve() if not output_dir.exists(): return [] if not output_dir.is_dir(): @@ -61,15 +63,24 @@ def list_experiments(system: pathlib.Path) -> list[tuple[str, pathlib.Path]]: def get_experiment(exp_id: str | pathlib.Path, system: str | pathlib.Path) -> cloudai.models.output.Experiment: - """Load an experiment by ID or by its result directory or JSON path.""" - if isinstance(exp_id, pathlib.Path): - path = exp_id.expanduser() - if path.suffix != ".json": - path /= "experiment.json" - else: - if exp_id in ("", ".", "..") or pathlib.Path(exp_id).name != exp_id: - raise ValueError(f"Invalid experiment ID: {exp_id!r}") - output_dir = Parser.parse_system(pathlib.Path(system).expanduser().resolve()).output_path.expanduser().resolve() - path = output_dir / exp_id / "experiment.json" + """Load an experiment by ID or result directory.""" + parsed_system = cloudai.core.Parser.parse_system(pathlib.Path(system).expanduser().resolve()) + system_output = parsed_system.output_path.expanduser().resolve() + + experiment_path = system_output / exp_id if isinstance(exp_id, str) else exp_id + experiment_path = experiment_path.expanduser().resolve() + experiment_file = experiment_path / "experiment.json" + if experiment_path.parent != system_output or not experiment_path.is_dir() or not experiment_file.is_file(): + raise ValueError(f"Invalid experiment ID: {experiment_path!r}") + + try: + experiment_data = experiment_file.read_text(encoding="utf-8") + except (OSError, ValueError) as error: + raise ValueError(f"Invalid experiment ID: {experiment_path!r}") from error + + try: + exp = cloudai.models.output.Experiment.model_validate_json(experiment_data) + except pydantic.ValidationError as error: + raise ValueError(f"Cannot read experiment: {experiment_path!r}") from error - return cloudai.models.output.Experiment.model_validate_json(path.read_text(encoding="utf-8")) + return exp diff --git a/tests/test_api.py b/tests/test_api.py index b1fd0218d..e46cb03e5 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -90,11 +90,14 @@ def test_run_and_list_experiments( ] assert cloudai.api.get_experiment(experiment.id, str(system)) == experiment assert cloudai.api.get_experiment(pathlib.Path(experiment.path), system) == experiment - assert cloudai.api.get_experiment(pathlib.Path(experiment.path) / "experiment.json", system) == experiment - with pytest.raises(FileNotFoundError): + with pytest.raises(ValueError, match="Invalid experiment ID"): + cloudai.api.get_experiment(pathlib.Path(experiment.path) / "experiment.json", system) + with pytest.raises(ValueError, match="Invalid experiment ID"): cloudai.api.get_experiment("missing", system) with pytest.raises(ValueError, match="Invalid experiment ID"): cloudai.api.get_experiment("../other", system) + with pytest.raises(ValueError, match="Cannot read experiment"): + cloudai.api.get_experiment(unparsed_path, system) assert signal.getsignal(signal.SIGINT) == signal_handler scenario.write_text('name = "invalid"\n') From f0ed6e583be4d4d2a49f915c52e9c96dd572f5af Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 15:56:03 +0200 Subject: [PATCH 06/13] Keep output directory override in CLI Signed-off-by: Ivan Podkidyshev --- src/cloudai/cli/cli.py | 3 ++- src/cloudai/handlers.py | 3 --- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index d68f5c958..ff3c31852 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -45,11 +45,12 @@ def handle_dry_run_and_run(args: argparse.Namespace) -> int: system, tests, scenario = cloudai.handlers.load_experiment( args.test_scenario, args.system_config, tests_dir=args.tests_dir, hook_dir=args.hook_dir ) + if args.output_dir is not None: + system.output_path = args.output_dir.absolute() runner = cloudai.handlers.create_experiment_runner( system, scenario, mode=args.mode, - output_dir=args.output_dir, single_sbatch=args.single_sbatch, ) except ( diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 1575e68c8..25e68d3fc 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -272,14 +272,11 @@ def create_experiment_runner( scenario: TestScenario, *, mode: str = "run", - output_dir: Path | None = None, single_sbatch: bool = False, ) -> Runner: """Prepare the system and create a runner for one experiment.""" if single_sbatch and not isinstance(system, SlurmSystem): raise ValueError("Single sbatch is only supported for Slurm systems.") - if output_dir is not None: - system.output_path = output_dir.absolute() if not prepare_output_dir(system.output_path): raise OSError(f"Cannot prepare output directory: {system.output_path}") if mode == "dry-run": From 2afa177f5a9e04553ccffd47d6a2fe974ab14532 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 18:29:11 +0200 Subject: [PATCH 07/13] Support run modes in Python API Signed-off-by: Ivan Podkidyshev --- doc/api.rst | 6 ++++-- src/cloudai/api.py | 9 ++++++++- tests/test_api.py | 32 ++++++++++++++++++-------------- 3 files changed, 30 insertions(+), 17 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index 138324cd1..e995caa70 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -20,8 +20,10 @@ Pass configuration file paths as ``pathlib.Path`` objects: ``run_experiment`` waits for the scenario to finish and returns an ``Experiment`` model with the same data saved in ``experiment.json``. Workload failures appear in -its status. Configuration and setup errors raise exceptions. Pass ``hook_dir`` to -use a custom hook directory. +its status. Configuration and setup errors raise exceptions. Use +``mode="dry-run"`` to generate commands without running workloads. Use +``single_sbatch=True`` for a single Slurm allocation. Pass ``hook_dir`` to use a +custom hook directory. ``list_experiments(system)`` returns ``(experiment_id, result_directory)`` pairs from the system's configured output directory. Directories without diff --git a/src/cloudai/api.py b/src/cloudai/api.py index de04e59cc..96a151f0d 100644 --- a/src/cloudai/api.py +++ b/src/cloudai/api.py @@ -27,6 +27,8 @@ def run_experiment( scenario: pathlib.Path, system: pathlib.Path, *, + mode: str = "run", + single_sbatch: bool = False, tests_dir: pathlib.Path | None = None, hook_dir: pathlib.Path | None = None, ) -> cloudai.models.output.Experiment: @@ -42,7 +44,12 @@ def run_experiment( tests_dir=tests_dir, hook_dir=hook_dir, ) - runner = cloudai.handlers.create_experiment_runner(parsed_system, parsed_scenario) + runner = cloudai.handlers.create_experiment_runner( + parsed_system, + parsed_scenario, + mode=mode, + single_sbatch=single_sbatch, + ) cloudai.handlers.execute_experiment(runner, tests) return runner.runner.experiment_output.snapshot() diff --git a/tests/test_api.py b/tests/test_api.py index e46cb03e5..de1f81488 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -45,22 +45,19 @@ def test_run_and_list_experiments( ) ) scenario = tmp_path / "scenario.toml" - scenario.write_text( - toml.dumps( + scenario_data = { + "name": "api-test", + "Tests": [ { - "name": "api-test", - "Tests": [ - { - "id": "sleep", - "name": "sleep", - "description": "API smoke test", - "test_template_name": "Sleep", - "cmd_args": {"seconds": 0}, - } - ], + "id": "sleep", + "name": "sleep", + "description": "API smoke test", + "test_template_name": "Sleep", + "cmd_args": {"seconds": 0}, } - ) - ) + ], + } + scenario.write_text(toml.dumps(scenario_data)) monkeypatch.setattr(StandaloneSystem, "is_job_completed", lambda *_: True) monkeypatch.setattr(StandaloneSystem, "is_job_running", lambda *_: False) monkeypatch.setattr( @@ -100,6 +97,13 @@ def test_run_and_list_experiments( cloudai.api.get_experiment(unparsed_path, system) assert signal.getsignal(signal.SIGINT) == signal_handler + scenario_data["name"] = "api-dry-run" + scenario.write_text(toml.dumps(scenario_data)) + dry_run_experiment = cloudai.api.run_experiment(scenario, system, mode="dry-run") + assert dry_run_experiment.tests[0].runs == [] + with pytest.raises(ValueError, match="Single sbatch is only supported for Slurm systems"): + cloudai.api.run_experiment(scenario, system, single_sbatch=True) + scenario.write_text('name = "invalid"\n') with pytest.raises(cloudai.core.TestScenarioParsingError): cloudai.api.run_experiment(scenario, system) From 19c315e516319debf19427f5a2f71f310c62a751 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 18:36:02 +0200 Subject: [PATCH 08/13] Split Python API tests by function Signed-off-by: Ivan Podkidyshev --- tests/test_api.py | 176 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 130 insertions(+), 46 deletions(-) diff --git a/tests/test_api.py b/tests/test_api.py index de1f81488..252dd0aa8 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -23,17 +23,17 @@ import cloudai.api import cloudai.core +import cloudai.models.output from cloudai.systems.standalone import StandaloneSystem from cloudai.systems.standalone.standalone_job import StandaloneJob from cloudai.systems.standalone.standalone_runner import StandaloneRunner from cloudai.workloads.sleep import SleepTestDefinition -def test_run_and_list_experiments( - standalone_system: StandaloneSystem, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch -) -> None: - system = tmp_path / "system.toml" - system.write_text( +@pytest.fixture +def system_config(standalone_system: StandaloneSystem, tmp_path: pathlib.Path) -> pathlib.Path: + path = tmp_path / "system.toml" + path.write_text( toml.dumps( { "name": standalone_system.name, @@ -44,20 +44,33 @@ def test_run_and_list_experiments( } ) ) - scenario = tmp_path / "scenario.toml" - scenario_data = { - "name": "api-test", - "Tests": [ + return path + + +@pytest.fixture +def scenario_config(tmp_path: pathlib.Path) -> pathlib.Path: + path = tmp_path / "scenario.toml" + path.write_text( + toml.dumps( { - "id": "sleep", - "name": "sleep", - "description": "API smoke test", - "test_template_name": "Sleep", - "cmd_args": {"seconds": 0}, + "name": "api-test", + "Tests": [ + { + "id": "sleep", + "name": "sleep", + "description": "API smoke test", + "test_template_name": "Sleep", + "cmd_args": {"seconds": 0}, + } + ], } - ], - } - scenario.write_text(toml.dumps(scenario_data)) + ) + ) + return path + + +@pytest.fixture +def mock_standalone_execution(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(StandaloneSystem, "is_job_completed", lambda *_: True) monkeypatch.setattr(StandaloneSystem, "is_job_running", lambda *_: False) monkeypatch.setattr( @@ -71,39 +84,110 @@ def test_run_and_list_experiments( lambda *_: cloudai.core.JobStatusResult(True, "test outcome"), ) + +@pytest.fixture +def saved_experiment( + standalone_system: StandaloneSystem, +) -> tuple[cloudai.models.output.Experiment, pathlib.Path]: + experiment_path = standalone_system.output_path / "experiment-1" + experiment_path.mkdir(parents=True) + experiment = cloudai.models.output.Experiment( + id=experiment_path.name, + name="experiment", + system_name=standalone_system.name, + status="completed", + path=str(experiment_path), + ) + (experiment_path / "experiment.json").write_text(experiment.model_dump_json()) + return experiment, experiment_path + + +@pytest.mark.parametrize(("mode", "expected_runs"), [("run", 1), ("dry-run", 0)]) +@pytest.mark.usefixtures("mock_standalone_execution") +def test_run_experiment( + mode: str, + expected_runs: int, + system_config: pathlib.Path, + scenario_config: pathlib.Path, +) -> None: signal_handler = signal.getsignal(signal.SIGINT) - assert cloudai.api.list_experiments(system) == [] - experiment = cloudai.api.run_experiment(scenario, system) + + experiment = cloudai.api.run_experiment(scenario_config, system_config, mode=mode) + assert experiment.status == "completed" - assert experiment.tests[0].runs[0].status == "completed" - experiment_path = pathlib.Path(experiment.path) - assert experiment.id == experiment_path.name - unparsed_path = standalone_system.output_path / "zzz-unparsed" - unparsed_path.mkdir() - (unparsed_path / "experiment.json").write_text("not JSON") - assert cloudai.api.list_experiments(system) == [ - (experiment.id, experiment_path), - (unparsed_path.name, unparsed_path.resolve()), - ] - assert cloudai.api.get_experiment(experiment.id, str(system)) == experiment - assert cloudai.api.get_experiment(pathlib.Path(experiment.path), system) == experiment - with pytest.raises(ValueError, match="Invalid experiment ID"): - cloudai.api.get_experiment(pathlib.Path(experiment.path) / "experiment.json", system) - with pytest.raises(ValueError, match="Invalid experiment ID"): - cloudai.api.get_experiment("missing", system) - with pytest.raises(ValueError, match="Invalid experiment ID"): - cloudai.api.get_experiment("../other", system) - with pytest.raises(ValueError, match="Cannot read experiment"): - cloudai.api.get_experiment(unparsed_path, system) + assert len(experiment.tests[0].runs) == expected_runs + assert experiment.id == pathlib.Path(experiment.path).name assert signal.getsignal(signal.SIGINT) == signal_handler - scenario_data["name"] = "api-dry-run" - scenario.write_text(toml.dumps(scenario_data)) - dry_run_experiment = cloudai.api.run_experiment(scenario, system, mode="dry-run") - assert dry_run_experiment.tests[0].runs == [] + +def test_run_experiment_rejects_single_sbatch_for_standalone( + system_config: pathlib.Path, + scenario_config: pathlib.Path, +) -> None: with pytest.raises(ValueError, match="Single sbatch is only supported for Slurm systems"): - cloudai.api.run_experiment(scenario, system, single_sbatch=True) + cloudai.api.run_experiment(scenario_config, system_config, single_sbatch=True) + + +def test_run_experiment_rejects_invalid_scenario( + system_config: pathlib.Path, + scenario_config: pathlib.Path, +) -> None: + scenario_config.write_text('name = "invalid"\n') - scenario.write_text('name = "invalid"\n') with pytest.raises(cloudai.core.TestScenarioParsingError): - cloudai.api.run_experiment(scenario, system) + cloudai.api.run_experiment(scenario_config, system_config) + + +def test_list_experiments(system_config: pathlib.Path, standalone_system: StandaloneSystem) -> None: + assert cloudai.api.list_experiments(system_config) == [] + experiment_paths = [standalone_system.output_path / name for name in ("experiment-1", "experiment-2")] + for path in experiment_paths: + path.mkdir(parents=True) + (path / "experiment.json").write_text("not parsed") + (standalone_system.output_path / "ignored").mkdir() + + assert cloudai.api.list_experiments(system_config) == [(path.name, path.resolve()) for path in experiment_paths] + + +@pytest.mark.parametrize("reference", ["id", "path"]) +def test_get_experiment( + reference: str, + system_config: pathlib.Path, + saved_experiment: tuple[cloudai.models.output.Experiment, pathlib.Path], +) -> None: + experiment, experiment_path = saved_experiment + exp_id = experiment.id if reference == "id" else experiment_path + + assert cloudai.api.get_experiment(exp_id, system_config) == experiment + + +@pytest.mark.parametrize( + ("reference", "message"), + [ + ("json-path", "Invalid experiment ID"), + ("missing", "Invalid experiment ID"), + ("traversal", "Invalid experiment ID"), + ("malformed", "Cannot read experiment"), + ], +) +def test_get_experiment_rejects_invalid_input( + reference: str, + message: str, + system_config: pathlib.Path, + saved_experiment: tuple[cloudai.models.output.Experiment, pathlib.Path], +) -> None: + _, experiment_path = saved_experiment + if reference == "json-path": + exp_id: str | pathlib.Path = experiment_path / "experiment.json" + elif reference == "missing": + exp_id = "missing" + elif reference == "traversal": + exp_id = "../other" + else: + malformed_path = experiment_path.parent / "malformed" + malformed_path.mkdir() + (malformed_path / "experiment.json").write_text("not JSON") + exp_id = malformed_path + + with pytest.raises(ValueError, match=message): + cloudai.api.get_experiment(exp_id, system_config) From 123bcb2bbe456824685d26d69ee75377ecf7c179 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 19:54:09 +0200 Subject: [PATCH 09/13] Report experiment startup through Python API Signed-off-by: Ivan Podkidyshev --- doc/api.rst | 4 +++- src/cloudai/api.py | 4 +++- src/cloudai/handlers.py | 10 +++++++++- tests/test_api.py | 21 +++++++++++++++++++-- 4 files changed, 34 insertions(+), 5 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index e995caa70..d3e07ca0d 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -23,7 +23,9 @@ model with the same data saved in ``experiment.json``. Workload failures appear its status. Configuration and setup errors raise exceptions. Use ``mode="dry-run"`` to generate commands without running workloads. Use ``single_sbatch=True`` for a single Slurm allocation. Pass ``hook_dir`` to use a -custom hook directory. +custom hook directory. Pass an ``on_start`` callback to receive the initial +``Experiment`` after ``experiment.json`` is created and before execution begins. +The function remains synchronous after calling the callback. ``list_experiments(system)`` returns ``(experiment_id, result_directory)`` pairs from the system's configured output directory. Directories without diff --git a/src/cloudai/api.py b/src/cloudai/api.py index 96a151f0d..b5615a463 100644 --- a/src/cloudai/api.py +++ b/src/cloudai/api.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import collections.abc import pathlib import pydantic @@ -31,6 +32,7 @@ def run_experiment( single_sbatch: bool = False, tests_dir: pathlib.Path | None = None, hook_dir: pathlib.Path | None = None, + on_start: collections.abc.Callable[[cloudai.models.output.Experiment], None] | None = None, ) -> cloudai.models.output.Experiment: """ Run a scenario synchronously and return its experiment snapshot. @@ -50,7 +52,7 @@ def run_experiment( mode=mode, single_sbatch=single_sbatch, ) - cloudai.handlers.execute_experiment(runner, tests) + cloudai.handlers.execute_experiment(runner, tests, on_start=on_start) return runner.runner.experiment_output.snapshot() diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 25e68d3fc..14b82a9c6 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -15,6 +15,7 @@ # limitations under the License. import argparse +import collections.abc import copy import logging import traceback @@ -26,6 +27,7 @@ import toml import yaml +import cloudai.models.output from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( BaseInstaller, @@ -287,7 +289,11 @@ def create_experiment_runner( def execute_experiment( - runner: Runner, tests: list[TestDefinition], *, enable_cache_without_check: bool = False + runner: Runner, + tests: list[TestDefinition], + *, + enable_cache_without_check: bool = False, + on_start: collections.abc.Callable[[cloudai.models.output.Experiment], None] | None = None, ) -> bool: """Install prerequisites, execute the scenario, and finalize its output.""" system = runner.runner.system @@ -296,6 +302,8 @@ def execute_experiment( successful = False try: runner.runner.experiment_output.write() + if on_start is not None: + on_start(runner.runner.experiment_output.snapshot()) logging.info(f"System Name: {system.name}") logging.info(f"Scheduler: {system.scheduler}") logging.info(f"Test Scenario Name: {scenario.name}") diff --git a/tests/test_api.py b/tests/test_api.py index 252dd0aa8..1ece7736b 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -111,9 +111,26 @@ def test_run_experiment( scenario_config: pathlib.Path, ) -> None: signal_handler = signal.getsignal(signal.SIGINT) + started_experiments: list[cloudai.models.output.Experiment] = [] + started_outputs: list[cloudai.models.output.Experiment] = [] + + def on_start(started: cloudai.models.output.Experiment) -> None: + started_experiments.append(started) + started_outputs.append(cloudai.api.get_experiment(started.id, system_config)) + + experiment = cloudai.api.run_experiment( + scenario_config, + system_config, + mode=mode, + on_start=on_start, + ) - experiment = cloudai.api.run_experiment(scenario_config, system_config, mode=mode) - + assert len(started_experiments) == 1 + assert started_experiments[0].status == "running" + assert started_experiments[0].id == experiment.id + assert started_outputs[0].status == "running" + assert started_outputs[0].id == experiment.id + assert started_outputs[0].tests[0].runs == [] assert experiment.status == "completed" assert len(experiment.tests[0].runs) == expected_runs assert experiment.id == pathlib.Path(experiment.path).name From 060a571a790878f8b1bb09ed440e50cead874b8e Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 20:49:16 +0200 Subject: [PATCH 10/13] Expose experiment cancellation on startup --- doc/api.rst | 5 +++-- src/cloudai/api.py | 4 +++- src/cloudai/handlers.py | 6 ++++-- tests/test_api.py | 21 ++++++++++++++++++--- 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index d3e07ca0d..c91585f9f 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -24,8 +24,9 @@ its status. Configuration and setup errors raise exceptions. Use ``mode="dry-run"`` to generate commands without running workloads. Use ``single_sbatch=True`` for a single Slurm allocation. Pass ``hook_dir`` to use a custom hook directory. Pass an ``on_start`` callback to receive the initial -``Experiment`` after ``experiment.json`` is created and before execution begins. -The function remains synchronous after calling the callback. +``Experiment`` and a cancellation function after ``experiment.json`` is created +and before execution begins. The function remains synchronous after calling the +callback. ``list_experiments(system)`` returns ``(experiment_id, result_directory)`` pairs from the system's configured output directory. Directories without diff --git a/src/cloudai/api.py b/src/cloudai/api.py index b5615a463..ce461bedd 100644 --- a/src/cloudai/api.py +++ b/src/cloudai/api.py @@ -32,7 +32,9 @@ def run_experiment( single_sbatch: bool = False, tests_dir: pathlib.Path | None = None, hook_dir: pathlib.Path | None = None, - on_start: collections.abc.Callable[[cloudai.models.output.Experiment], None] | None = None, + on_start: ( + collections.abc.Callable[[cloudai.models.output.Experiment, collections.abc.Callable[[], None]], None] | None + ) = None, ) -> cloudai.models.output.Experiment: """ Run a scenario synchronously and return its experiment snapshot. diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 14b82a9c6..f89418ab7 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -293,7 +293,9 @@ def execute_experiment( tests: list[TestDefinition], *, enable_cache_without_check: bool = False, - on_start: collections.abc.Callable[[cloudai.models.output.Experiment], None] | None = None, + on_start: ( + collections.abc.Callable[[cloudai.models.output.Experiment, collections.abc.Callable[[], None]], None] | None + ) = None, ) -> bool: """Install prerequisites, execute the scenario, and finalize its output.""" system = runner.runner.system @@ -303,7 +305,7 @@ def execute_experiment( try: runner.runner.experiment_output.write() if on_start is not None: - on_start(runner.runner.experiment_output.snapshot()) + on_start(runner.runner.experiment_output.snapshot(), runner.runner.shutdown) logging.info(f"System Name: {system.name}") logging.info(f"Scheduler: {system.scheduler}") logging.info(f"Test Scenario Name: {scenario.name}") diff --git a/tests/test_api.py b/tests/test_api.py index 1ece7736b..72714d18b 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import collections.abc import datetime import pathlib import signal @@ -102,10 +103,19 @@ def saved_experiment( return experiment, experiment_path -@pytest.mark.parametrize(("mode", "expected_runs"), [("run", 1), ("dry-run", 0)]) +@pytest.mark.parametrize( + ("mode", "cancel_on_start", "expected_status", "expected_runs"), + [ + ("run", False, "completed", 1), + ("dry-run", False, "completed", 0), + ("run", True, "cancelled", 0), + ], +) @pytest.mark.usefixtures("mock_standalone_execution") def test_run_experiment( mode: str, + cancel_on_start: bool, + expected_status: str, expected_runs: int, system_config: pathlib.Path, scenario_config: pathlib.Path, @@ -114,9 +124,14 @@ def test_run_experiment( started_experiments: list[cloudai.models.output.Experiment] = [] started_outputs: list[cloudai.models.output.Experiment] = [] - def on_start(started: cloudai.models.output.Experiment) -> None: + def on_start( + started: cloudai.models.output.Experiment, + cancel: collections.abc.Callable[[], None], + ) -> None: started_experiments.append(started) started_outputs.append(cloudai.api.get_experiment(started.id, system_config)) + if cancel_on_start: + cancel() experiment = cloudai.api.run_experiment( scenario_config, @@ -131,7 +146,7 @@ def on_start(started: cloudai.models.output.Experiment) -> None: assert started_outputs[0].status == "running" assert started_outputs[0].id == experiment.id assert started_outputs[0].tests[0].runs == [] - assert experiment.status == "completed" + assert experiment.status == expected_status assert len(experiment.tests[0].runs) == expected_runs assert experiment.id == pathlib.Path(experiment.path).name assert signal.getsignal(signal.SIGINT) == signal_handler From 870a0744c9c828709dbf60799e3cebc2c28e18af Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Thu, 1 Oct 2026 21:01:08 +0200 Subject: [PATCH 11/13] Register CLI signals after experiment startup --- src/cloudai/cli/cli.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index ff3c31852..36e685837 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -18,6 +18,7 @@ import logging import logging.config import signal +import types from pathlib import Path from typing import Callable @@ -65,10 +66,19 @@ def handle_dry_run_and_run(args: argparse.Namespace) -> int: logging.error(str(exc)) return 1 - register_signal_handlers(runner.cancel_on_signal) + def on_start(_experiment: object, cancel: Callable[[], None]) -> None: + def cancel_on_signal(signum: int, _frame: types.FrameType | None) -> None: + logging.info(f"Signal {signum} received, shutting down...") + cancel() + + register_signal_handlers(cancel_on_signal) + try: successful = cloudai.handlers.execute_experiment( - runner, tests, enable_cache_without_check=args.enable_cache_without_check + runner, + tests, + enable_cache_without_check=args.enable_cache_without_check, + on_start=on_start, ) except cloudai.handlers.InstallationError: return 1 From c5267683a82a8bc5037973ee9c2c1ad19c9a3f3a Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 6 Oct 2026 13:27:10 +0200 Subject: [PATCH 12/13] Use consistent handlers imports in CLI --- src/cloudai/cli/cli.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 36e685837..67f959574 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -26,12 +26,6 @@ import cloudai.core import cloudai.handlers -from cloudai.handlers import ( - handle_generate_report, - handle_install_and_uninstall, - handle_list_registered_items, - handle_verify_all_configs, -) def register_signal_handlers(signal_handler: Callable) -> None: @@ -211,7 +205,7 @@ def install(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_dir: Pat hook_dir=hook_dir, mode="install", ) - exit(handle_install_and_uninstall(args)) + exit(cloudai.handlers.handle_install_and_uninstall(args)) @main.command() @@ -225,7 +219,7 @@ def uninstall(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_dir: P hook_dir=hook_dir, mode="uninstall", ) - exit(handle_install_and_uninstall(args)) + exit(cloudai.handlers.handle_install_and_uninstall(args)) @main.command() @@ -310,7 +304,7 @@ def generate_report(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_ hook_dir=hook_dir, result_dir=result_dir, ) - exit(handle_generate_report(args)) + exit(cloudai.handlers.handle_generate_report(args)) @main.command() @@ -326,7 +320,7 @@ def generate_report(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_ def verify_configs(configs_dir: Path, tests_dir: Path): """Verify the configuration TOML files.""" args = argparse.Namespace(configs_dir=configs_dir, tests_dir=tests_dir) - raise click.exceptions.Exit(1 if handle_verify_all_configs(args) else 0) + raise click.exceptions.Exit(1 if cloudai.handlers.handle_verify_all_configs(args) else 0) @main.command() @@ -334,4 +328,4 @@ def verify_configs(configs_dir: Path, tests_dir: Path): @click.option("-v", "--verbose", is_flag=True, default=False, help="Verbose output.") def list(type: str, verbose: bool): """List available in Registry items.""" - handle_list_registered_items(type, verbose) + cloudai.handlers.handle_list_registered_items(type, verbose) From c9513e6b1f4b35108cc28ff8c122c2f2f9901a4d Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 6 Oct 2026 13:36:35 +0200 Subject: [PATCH 13/13] Allow explicit runner class overrides --- src/cloudai/_core/runner.py | 11 ++++++----- tests/test_handlers.py | 9 +++++++++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/src/cloudai/_core/runner.py b/src/cloudai/_core/runner.py index d2b150ae0..3066895a4 100644 --- a/src/cloudai/_core/runner.py +++ b/src/cloudai/_core/runner.py @@ -90,11 +90,12 @@ def create_runner( """ registry = Registry() scheduler_type = system.scheduler.lower() - if scheduler_type not in registry.runners_map: - msg = f"No runner registered for scheduler: {scheduler_type}" - logging.error(msg) - raise NotImplementedError(msg) - runner_class = runner_class or registry.runners_map[scheduler_type] + if runner_class is None: + if scheduler_type not in registry.runners_map: + msg = f"No runner registered for scheduler: {scheduler_type}" + logging.error(msg) + raise NotImplementedError(msg) + runner_class = registry.runners_map[scheduler_type] logging.info(f"Creating {runner_class.__name__}") if not system.output_path.exists(): diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 71ae9fff1..ea5a09f64 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -102,6 +102,15 @@ def test_runner_warns_about_legacy_create_runner_signature(slurm_system: SlurmSy Runner(mode="dry-run", system=slurm_system, test_scenario=scenario) +def test_runner_uses_explicit_class_without_registered_scheduler(slurm_system: SlurmSystem, dse_tr: TestRun) -> None: + slurm_system.scheduler = "custom" + scenario = TestScenario(name="test_scenario", test_runs=[dse_tr]) + + runner = Runner(mode="dry-run", system=slurm_system, test_scenario=scenario, runner_class=SlurmRunner) + + assert isinstance(runner.runner, SlurmRunner) + + @pytest.mark.parametrize("dep", ["start_post_comp", "start_post_init", "end_post_comp"]) def test_dse_run_does_not_support_dependencies( slurm_system: SlurmSystem, dse_tr: TestRun, dep: str, caplog: pytest.LogCaptureFixture