diff --git a/doc/api.rst b/doc/api.rst new file mode 100644 index 000000000..c91585f9f --- /dev/null +++ b/doc/api.rst @@ -0,0 +1,39 @@ +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. 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`` 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 +``experiment.json`` are skipped. A missing output directory gives an empty list; +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 to load it +directly. Invalid IDs and missing, unreadable, or malformed experiment files +raise ``ValueError``. 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..3066895a4 100644 --- a/src/cloudai/_core/runner.py +++ b/src/cloudai/_core/runner.py @@ -16,6 +16,7 @@ import datetime import logging +import warnings from types import FrameType from typing import Optional @@ -44,11 +45,34 @@ 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, + ): 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: + 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) + + def create_runner( + self, + mode: str, + system: System, + test_scenario: TestScenario, + *, + runner_class: type[BaseRunner] | None = None, + ) -> BaseRunner: """ Dynamically create a runner instance based on the system's scheduler type. @@ -56,6 +80,7 @@ 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. Returns: BaseRunner: A runner instance suitable for the system. @@ -65,11 +90,12 @@ def create_runner(self, mode: str, system: System, test_scenario: TestScenario) """ 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 = 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/src/cloudai/api.py b/src/cloudai/api.py new file mode 100644 index 000000000..ce461bedd --- /dev/null +++ b/src/cloudai/api.py @@ -0,0 +1,97 @@ +# 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 collections.abc +import pathlib + +import pydantic + +import cloudai.core +import cloudai.handlers +import cloudai.models.output + + +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, + 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. + + 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, + mode=mode, + single_sbatch=single_sbatch, + ) + cloudai.handlers.execute_experiment(runner, tests, on_start=on_start) + 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 = cloudai.core.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) + + 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: + """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 exp diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 66d478742..67f959574 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -17,17 +17,69 @@ import argparse import logging import logging.config +import signal +import types from pathlib import Path +from typing import Callable import click -from cloudai.handlers import ( - handle_dry_run_and_run, - handle_generate_report, - handle_install_and_uninstall, - handle_list_registered_items, - handle_verify_all_configs, -) +import cloudai.core +import cloudai.handlers + + +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 + ) + 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, + 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 + + 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, + on_start=on_start, + ) + 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: @@ -153,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() @@ -167,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() @@ -252,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() @@ -268,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() @@ -276,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) diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index de41076f9..6d5b105df 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -14,36 +14,41 @@ # 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( - 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, ) - 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 +61,25 @@ 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. " + "This compatibility shim will be removed in CloudAI 1.9.1.", + 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. " + "This compatibility shim will be removed in CloudAI 1.9.1.", + DeprecationWarning, + stacklevel=2, + ) + return cloudai.handlers.handle_non_dse_job(runner) diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 09402d619..f89418ab7 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -15,25 +15,24 @@ # limitations under the License. import argparse +import collections.abc import copy import logging -import signal import traceback 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 import yaml +import cloudai.models.output from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( BaseInstaller, CloudAIGymEnv, Installable, - InstallStatusResult, - MissingTestError, Parser, Registry, Runner, @@ -130,7 +129,7 @@ def _scenario_installables(scenario: TestScenario) -> list[Installable]: return installables -def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: +def handle_dse_job(runner: Runner) -> int: registry = Registry() original_test_runs = copy.deepcopy(runner.runner.test_scenario.test_runs) @@ -182,7 +181,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 runner.runner.mode == "run": runner.runner.test_scenario.test_runs = original_test_runs generate_reports( runner.runner.system, @@ -243,135 +242,102 @@ 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) -> 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: - """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) +class InstallationError(RuntimeError): + """Workload prerequisites could not be installed.""" -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 +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 - if args.output_dir: - system.output_path = args.output_dir.absolute() +def create_experiment_runner( + system: System, + scenario: TestScenario, + *, + mode: str = "run", + 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 not prepare_output_dir(system.output_path): - return None - - if args.mode == "dry-run": + raise OSError(f"Cannot prepare output directory: {system.output_path}") + if 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}") - + 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( + runner: Runner, + tests: list[TestDefinition], + *, + enable_cache_without_check: bool = False, + 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 + scenario = runner.runner.test_scenario + mode = runner.runner.mode 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 + if on_start is not None: + 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}") + 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 = handle_dse_job(runner) == 0 + else: + logging.error("Mixing DSE and non-DSE jobs is not allowed.") + return successful finally: runner.runner.finish_output(successful) 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..3857169cd 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..72714d18b --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,225 @@ +# 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 collections.abc +import datetime +import pathlib +import signal + +import pytest +import toml + +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 + + +@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, + "scheduler": "standalone", + "install_path": str(standalone_system.install_path), + "output_path": str(standalone_system.output_path), + "monitor_interval": 0, + } + ) + ) + return path + + +@pytest.fixture +def scenario_config(tmp_path: pathlib.Path) -> pathlib.Path: + path = tmp_path / "scenario.toml" + path.write_text( + toml.dumps( + { + "name": "api-test", + "Tests": [ + { + "id": "sleep", + "name": "sleep", + "description": "API smoke test", + "test_template_name": "Sleep", + "cmd_args": {"seconds": 0}, + } + ], + } + ) + ) + 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( + 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"), + ) + + +@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", "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, +) -> 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, + 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, + system_config, + mode=mode, + on_start=on_start, + ) + + 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 == 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 + + +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_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') + + with pytest.raises(cloudai.core.TestScenarioParsingError): + 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) 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 1e78d117f..ea5a09f64 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 @@ -50,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): @@ -97,6 +96,21 @@ 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) + + +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 @@ -113,8 +127,8 @@ 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 + 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 @@ -158,9 +172,9 @@ 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, 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 +269,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, @@ -393,11 +407,11 @@ 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) - 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 @@ -411,9 +425,9 @@ 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, argparse.Namespace(mode="dry-run")) == 1 + assert handle_dse_job(runner) == 1 assert CustomRunStubAgent.run_calls == 1 @@ -433,9 +447,9 @@ 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, argparse.Namespace(mode="dry-run")) == 1 + assert handle_dse_job(runner) == 1 assert CustomRunStubAgent.run_calls == 2 @@ -453,12 +467,12 @@ 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) 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 @@ -474,10 +488,10 @@ 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, argparse.Namespace(mode="dry-run")) + handle_dse_job(runner) assert CustomRunStubAgent.run_calls == 1 @@ -491,11 +505,11 @@ 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"): - handle_dse_job(runner, argparse.Namespace(mode="run")) + handle_dse_job(runner) failure_report = tmp_path / "dse_failure.txt" assert failure_report.exists()