diff --git a/doc/api.rst b/doc/api.rst new file mode 100644 index 000000000..f65555d7c --- /dev/null +++ b/doc/api.rst @@ -0,0 +1,76 @@ +Python API +========== + +Use ``cloudai.api`` to validate scenarios, run experiments, and read their results +from Python. The CLI and API share the same execution logic. + +Run an experiment +----------------- + +Pass configuration files as ``pathlib.Path`` objects: + +.. code-block:: python + + from pathlib import Path + + import cloudai.api + + scenario = Path("conf/common/test_scenario/sleep.toml") + system = Path("conf/common/system/example_slurm_cluster.toml") + tests_dir = Path("conf/common/test") + + valid, errors = cloudai.api.validate_scenario(scenario, system, tests_dir=tests_dir) + if not valid: + raise ValueError(errors) + + experiment = cloudai.api.run_experiment(scenario, system, tests_dir=tests_dir) + print(experiment.status, experiment.path) + +Test ``path`` references resolve relative to the scenario file. Other relative +paths keep their usual meaning relative to the working directory. + +``run_experiment`` installs workload prerequisites as needed and returns an +``Experiment`` model containing the saved results. Each invocation gets its own +result directory. Use ``experiment.model_dump()`` or ``experiment.model_dump_json()`` +to serialize the result. + +Optional keyword arguments: + +* ``tests_dir`` and ``hook_dir`` select test definitions and hooks. +* ``output_dir`` overrides the system's results directory. +* ``dry_run=True`` generates workload commands without submitting jobs. +* ``single_sbatch=True`` runs a Slurm scenario in one allocation. +* ``enable_cache_without_check=True`` uses installed workload components without + checking them first, as with the CLI option. + +Find and read results +--------------------- + +``list_experiments(system)`` returns ``(experiment_id, result_directory)`` pairs +from the system's ``output_path``, sorted by directory name. Directories without +``experiment.json`` are skipped. A missing results directory returns an empty list. +When using ``output_dir`` for a run, use that directory as ``output_path`` in the +system configuration passed to ``list_experiments``. + +``get_experiment(exp)`` accepts a result directory or the ``experiment.json`` file +itself, as either ``str`` or ``Path``. It reads the saved snapshot without querying +the scheduler. See :doc:`reporting` for the result fields. + +Errors +------ + +``validate_scenario`` returns ``(True, {})`` for valid inputs, or ``(False, errors)`` +with ``system`` or ``scenario`` keys and error messages. It checks configuration and +execution constraints without installing workloads, creating a results directory, +or querying cluster availability. + +``run_experiment`` raises exceptions for invalid configuration and setup errors. +Workload failures appear in the experiment and run statuses. + +Missing result files raise ``FileNotFoundError``. Invalid experiment JSON raises +``pydantic.ValidationError``. Listing also raises for unreadable or invalid result +files rather than silently omitting them. + +The API does not install signal handlers or invoke the CLI's logging configuration. +The legacy public functions in ``cloudai.cli.handlers`` remain available with +deprecation warnings; use ``cloudai.api`` for new integrations. diff --git a/doc/index.rst b/doc/index.rst index 48fb1332d..6b951d7ba 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -14,6 +14,7 @@ This document contains the following chapters: Tutorial USER_GUIDE reporting + api systems workloads/index DEV @@ -22,6 +23,7 @@ This document contains the following chapters: - :doc:`Tutorial` - :doc:`USER_GUIDE` - :doc:`reporting` +- :doc:`api` - :doc:`systems` - :doc:`workloads/index` - :doc:`DEV` diff --git a/pyproject.toml b/pyproject.toml index cb643ee67..0dcfab3d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,14 +153,20 @@ root_package = "cloudai" "cloudai.util", "cloudai.cli", "cloudai.handlers", + "cloudai.api", "cloudai.report_generator", ] [[tool.importlinter.contracts]] name = "Report generator is leaf dependency" type = "forbidden" - forbidden_modules = ["cloudai.workloads", "cloudai.cli", "cloudai.handlers"] - allow_indirect_imports = true # allow "from cloudai.core import ..." + forbidden_modules = [ + "cloudai.workloads", + "cloudai.cli", + "cloudai.handlers", + "cloudai.api", + ] + allow_indirect_imports = true # allow "from cloudai.core import ..." source_modules = ["cloudai.report_generator"] [[tool.importlinter.contracts]] @@ -171,6 +177,7 @@ root_package = "cloudai" "cloudai.workloads", "cloudai.cli", "cloudai.handlers", + "cloudai.api", ] allow_indirect_imports = true source_modules = ["cloudai.util"] diff --git a/src/cloudai/_core/base_runner.py b/src/cloudai/_core/base_runner.py index 681611c8b..9af6b6408 100644 --- a/src/cloudai/_core/base_runner.py +++ b/src/cloudai/_core/base_runner.py @@ -161,7 +161,7 @@ def submit_test(self, tr: TestRun): self.update_run_output(job) except JobSubmissionError as e: logging.error(e) - exit(1) + raise def on_job_submit(self, tr: TestRun) -> None: return diff --git a/src/cloudai/_core/runner.py b/src/cloudai/_core/runner.py index ee636dd30..98e8a2e5e 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: Runner override for this invocation. + output_path: Exact scenario result directory, when already 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..5cfaed099 --- /dev/null +++ b/src/cloudai/api.py @@ -0,0 +1,96 @@ +# 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. + +from pathlib import Path + +import cloudai.handlers +import cloudai.models.output + + +def run_experiment( + scenario: Path, + system: Path, + *, + tests_dir: Path | None = None, + hook_dir: Path | None = None, + output_dir: Path | None = None, + dry_run: bool = False, + single_sbatch: bool = False, + enable_cache_without_check: bool = False, +) -> cloudai.models.output.Experiment: + """ + Run a scenario and return its experiment snapshot. + + Pass configuration files as Path. Relative test paths resolve relative to + the scenario file. Optional directories have the same meaning as in the CLI. + The call waits for execution to finish. + + Configuration and setup errors raise exceptions. Workload failures are + recorded in the returned experiment's status. This function does not invoke + CLI logging configuration or install signal handlers. + """ + return cloudai.handlers.run_experiment( + scenario, + system, + tests_dir=tests_dir, + hook_dir=hook_dir, + output_dir=output_dir, + dry_run=dry_run, + single_sbatch=single_sbatch, + enable_cache_without_check=enable_cache_without_check, + ) + + +def list_experiments(system: Path) -> list[tuple[str, Path]]: + """ + Return experiment IDs and local result directories under system.output_path. + + Pass the system configuration file as Path. A missing results + directory returns an empty list. Unreadable or invalid experiment files raise + exceptions rather than returning an incomplete catalog. + """ + return cloudai.handlers.list_experiments(system) + + +def validate_scenario( + scenario: Path, + system: Path, + *, + tests_dir: Path | None = None, + hook_dir: Path | None = None, + single_sbatch: bool = False, +) -> tuple[bool, dict[str, str]]: + """ + Validate configuration files without installing or running jobs. + + Return (True, {}) on success, otherwise (False, errors), with configuration + names as keys and error messages as values. This checks configuration and + execution constraints, not whether cluster resources are currently available. + """ + return cloudai.handlers.validate_scenario( + scenario, system, tests_dir=tests_dir, hook_dir=hook_dir, single_sbatch=single_sbatch + ) + + +def get_experiment(exp: str | Path) -> cloudai.models.output.Experiment: + """ + Read experiment.json from a result directory or a JSON file path. + + Both str and Path represent filesystem paths here. Missing files raise + FileNotFoundError; invalid JSON or schema raises pydantic.ValidationError. + This reads the saved snapshot without querying the scheduler. + """ + return cloudai.handlers.get_experiment(exp) diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index b2c920a95..1396650b6 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -30,22 +30,19 @@ import cloudai.handlers from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( - InstallStatusResult, + JobSubmissionError, MissingTestError, Parser, Registry, - Runner, System, - TestDefinition, + SystemConfigParsingError, + TestConfigParsingError, TestParser, - TestScenario, TestScenarioParsingError, ) from cloudai.parser import HOOK_ROOT -from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem from cloudai.test_parser import load_test_toml_file from cloudai.toml_utils import format_toml_decode_error -from cloudai.util import prepare_output_dir def setup_logging(log_file: str, log_level: str) -> None: @@ -170,7 +167,8 @@ def handle_generate_report(args: argparse.Namespace) -> int: """ parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) system, _, test_scenario = parser.parse(args.tests_dir, args.test_scenario) - assert test_scenario is not None + if test_scenario is None: + raise ValueError("A test scenario is required.") cloudai.handlers.generate_reports(system, test_scenario, args.result_dir) @@ -452,118 +450,31 @@ def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: return 0 -def _setup_system_and_scenario( - args: argparse.Namespace, -) -> tuple[System, TestScenario, list[TestDefinition]] | None: - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - try: - system, tests, test_scenario = parser.parse(args.tests_dir, args.test_scenario) - except MissingTestError as e: - logging.error(e.message) - return None - - assert test_scenario is not None - - if args.output_dir: - system.output_path = args.output_dir.absolute() - - if not prepare_output_dir(system.output_path): - return None - - if args.mode == "dry-run": - system.monitor_interval = 1 - system.update() - - return system, test_scenario, tests - - -def _handle_single_sbatch(args: argparse.Namespace, system: System) -> bool: - if not args.single_sbatch: - return True - - if not isinstance(system, SlurmSystem): - logging.error("Single sbatch is only supported for Slurm systems.") - return False - - Registry().update_runner("slurm", SingleSbatchRunner) - return True - - -def _check_installation( - args: argparse.Namespace, system: System, tests: list[TestDefinition], test_scenario: TestScenario -) -> InstallStatusResult: - logging.info("Checking if workloads components are installed.") - installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) - - if args.enable_cache_without_check: - result = installer.mark_as_installed(installables) - else: - result = installer.is_installed(installables) - - return result - - def handle_dry_run_and_run(args: argparse.Namespace) -> int: - setup_result = _setup_system_and_scenario(args) - if setup_result is None: - return 1 - system, test_scenario, tests = setup_result - try: - validate_domain_randomization_active(test_scenario) - except TestScenarioParsingError as e: - logging.error(str(e)) - return 1 - - if not _handle_single_sbatch(args, system): - return 1 - - logging.info(f"System Name: {system.name}") - logging.info(f"Scheduler: {system.scheduler}") - logging.info(f"Test Scenario Name: {test_scenario.name}") - - result = _check_installation(args, system, tests, test_scenario) - if args.mode == "run" and not result.success: - logging.info("Not all workloads components are installed. Installing...") - installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) - - result = installer.install(installables) - if result.success: - _log_installation_dirs("CloudAI is successfully installed into", system) - else: - logging.error("Failed to install workloads components.") - logging.error(result.message) - return 1 - elif args.mode == "dry-run": - # simulate installation for dry-run - installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) - result = installer.mark_as_installed(installables) - if not result.success: - logging.warning("Failed to mark workloads components as installed for dry-run.") - - logging.info(test_scenario.pretty_print()) - - runner = Runner(args.mode, system, test_scenario) - register_signal_handlers(runner.cancel_on_signal) - logging.info(f"Scenario results will be stored at: {runner.runner.scenario_root}") - - successful = False - try: - runner.runner.experiment_output.write() - has_dse = any(tr.is_dse_job for tr in test_scenario.test_runs) - if args.single_sbatch or not has_dse: # in this mode cases are unrolled using grid search - successful = cloudai.handlers.handle_non_dse_job(runner, args) - return 0 - - if all(tr.is_dse_job for tr in test_scenario.test_runs): - result = cloudai.handlers.handle_dse_job(runner, args) - successful = result == 0 - return result - - logging.error("Mixing DSE and non-DSE jobs is not allowed.") + system, tests, scenario = cloudai.handlers.load_scenario( + 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, dry_run=args.mode == "dry-run", single_sbatch=args.single_sbatch + ) + register_signal_handlers(runner.cancel_on_signal) + logging.info("Results directory: %s", runner.runner.scenario_root) + return cloudai.handlers.execute_experiment(runner, tests, args.enable_cache_without_check) + except ( + MissingTestError, + SystemConfigParsingError, + TestConfigParsingError, + TestScenarioParsingError, + JobSubmissionError, + OSError, + ValueError, + RuntimeError, + ) as exc: + logging.error(str(exc)) return 1 - finally: - runner.runner.finish_output(successful) def common_options(f): diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 7e27d172d..be76c6555 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -14,12 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import argparse import functools import typing import warnings import cloudai.cli.cli import cloudai.configurator.env_params +import cloudai.core import cloudai.handlers import cloudai.test_parser import cloudai.toml_utils @@ -60,5 +62,12 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: format_toml_decode_error = _deprecated(cloudai.toml_utils.format_toml_decode_error) prepare_output_dir = _deprecated(cloudai.util.prepare_output_dir) -handle_dse_job = _deprecated(cloudai.handlers.handle_dse_job) -handle_non_dse_job = _deprecated(cloudai.handlers.handle_non_dse_job) + +@_deprecated +def handle_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> int: + return cloudai.handlers.handle_dse_job(runner, args.mode) + + +@_deprecated +def handle_non_dse_job(runner: cloudai.core.Runner, args: argparse.Namespace) -> bool: + return cloudai.handlers.handle_non_dse_job(runner) diff --git a/src/cloudai/core.py b/src/cloudai/core.py index 2ebd13c8c..d7eab1cd9 100644 --- a/src/cloudai/core.py +++ b/src/cloudai/core.py @@ -25,6 +25,7 @@ from ._core.exceptions import ( JobFailureError, JobIdRetrievalError, + JobSubmissionError, MissingTestError, SystemConfigParsingError, TestConfigParsingError, @@ -102,6 +103,7 @@ "JobFailureError", "JobIdRetrievalError", "JobStatusResult", + "JobSubmissionError", "JsonGenStrategy", "MetricErrorSentinel", "MetricValue", diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py index 2563e6e86..4995b8264 100644 --- a/src/cloudai/handlers.py +++ b/src/cloudai/handlers.py @@ -14,25 +14,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse import copy +import datetime import logging +import tempfile import traceback from pathlib import Path from typing import Optional +import cloudai.models.output +import cloudai.output +from cloudai.configurator.env_params import validate_domain_randomization_active from cloudai.core import ( BaseInstaller, CloudAIGymEnv, Installable, + MissingTestError, + Parser, Registry, Runner, System, + SystemConfigParsingError, + TestConfigParsingError, TestScenario, + TestScenarioParsingError, ) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import TestDefinition -from cloudai.systems.slurm import SlurmSystem +from cloudai.parser import HOOK_ROOT +from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem +from cloudai.util import prepare_output_dir def prepare_installation( @@ -66,7 +77,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, mode: str) -> int: registry = Registry() original_test_runs = copy.deepcopy(runner.runner.test_scenario.test_runs) @@ -118,7 +129,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, @@ -179,40 +190,168 @@ 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 handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: C901 - registry = Registry() - if item_type.lower() == "reports": - print("Available scenario reports:") - for idx, (name, report) in enumerate(sorted(registry.scenario_reports.items()), start=1): - string = f'{idx}. "{name}" {report.__name__}' - if verbose: - string += f" (config={registry.report_configs[name].model_dump_json(indent=None)})" - print(string) - elif item_type.lower() == "agents": - print("Available agents:") - for idx, name in enumerate(registry.agent_names(), start=1): - agent = registry.get_agent(name) - string = f'{idx}. "{name}" class={agent.__name__}' - if verbose: - string += f"{agent.__doc__}" - print(string) - elif item_type.lower() == "reward-functions": - print("Available reward functions:") - for idx, name in enumerate(registry.reward_function_names(), start=1): - reward_function = registry.get_reward_function(name) - callable_name = getattr(reward_function, "__name__", type(reward_function).__name__) - string = f'{idx}. "{name}" function={callable_name}' - if verbose: - documentation = getattr(reward_function, "__doc__", None) - if documentation: - string += f" {documentation}" - print(string) - - return 0 +def load_scenario( + scenario: Path, system: Path, *, tests_dir: Path | None = None, hook_dir: Path | None = None +) -> tuple[System, list[TestDefinition], TestScenario]: + """Parse an experiment without installation, scheduler queries, or process exits.""" + try: + system.stat() + except OSError as exc: + raise SystemConfigParsingError(str(exc)) from exc + parser = Parser(system, hook_dir or HOOK_ROOT, exit_on_error=False) + parsed_system, tests, parsed_scenario = parser.parse(tests_dir, scenario) + if parsed_scenario is None: + raise TestScenarioParsingError("A test scenario is required.") + return parsed_system, tests, parsed_scenario + + +def validate_experiment(system: System, scenario: TestScenario, single_sbatch: bool = False) -> None: + """Check execution constraints shared by validation and execution.""" + validate_domain_randomization_active(scenario) + if single_sbatch: + if not isinstance(system, SlurmSystem): + raise TestScenarioParsingError("Single sbatch is only supported for Slurm systems.") + return + dse_runs = [tr for tr in scenario.test_runs if tr.is_dse_job] + if dse_runs and len(dse_runs) != len(scenario.test_runs): + raise TestScenarioParsingError("Mixing DSE and non-DSE jobs is not allowed.") + for tr in dse_runs: + if tr.dependencies: + raise TestScenarioParsingError("Dependencies are not supported for DSE jobs.") + registry = Registry() + if not registry.has_agent(tr.test.agent): + raise TestScenarioParsingError(f"No agent available for type: {tr.test.agent}.") + registry.get_agent(tr.test.agent).get_config_class()(**(tr.test.agent_config or {})) + + +def create_experiment_runner( + system: System, + scenario: TestScenario, + *, + dry_run: bool = False, + single_sbatch: bool = False, + result_dir: Path | None = None, +) -> Runner: + """Create a runner without changing the registered defaults.""" + validate_experiment(system, scenario, single_sbatch) + if prepare_output_dir(system.output_path) is None: + raise OSError(f"Cannot prepare output directory: {system.output_path}") + if dry_run: + system.monitor_interval = 1 + return Runner( + "dry-run" if dry_run else "run", + 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) -> int: + """Install prerequisites, execute the scenario, and finalize its output.""" + system = runner.runner.system + scenario = runner.runner.test_scenario + successful = False + try: + runner.runner.experiment_output.write() + system.update() + logging.info("System Name: %s", system.name) + logging.info("Scheduler: %s", system.scheduler) + logging.info("Test Scenario Name: %s", scenario.name) + installables, installer = prepare_installation(system, tests, scenario) + if enable_cache_without_check or runner.runner.mode == "dry-run": + result = installer.mark_as_installed(installables) + else: + result = installer.is_installed(installables) + if runner.runner.mode == "run" and not result.success: + logging.info("Not all workload components are installed. Installing...") + result = installer.install(installables) + if not result.success: + raise RuntimeError(f"Failed to install workload components: {result.message}") + elif runner.runner.mode == "dry-run" and not result.success: + logging.warning("Failed to mark workload components as installed for dry-run.") + logging.info(scenario.pretty_print()) + if isinstance(runner.runner, SingleSbatchRunner) or not any(tr.is_dse_job for tr in scenario.test_runs): + successful = handle_non_dse_job(runner) + return 0 + result_code = handle_dse_job(runner, runner.runner.mode) + successful = result_code == 0 + return result_code + finally: + runner.runner.finish_output(successful) + + +def run_experiment( + scenario: Path, + system: Path, + *, + tests_dir: Path | None = None, + hook_dir: Path | None = None, + output_dir: Path | None = None, + dry_run: bool = False, + single_sbatch: bool = False, + enable_cache_without_check: bool = False, +) -> cloudai.models.output.Experiment: + """Run an experiment from configuration files and return its results.""" + parsed_system, tests, parsed_scenario = load_scenario( + scenario.expanduser().resolve(), system.expanduser().resolve(), tests_dir=tests_dir, hook_dir=hook_dir + ) + validate_experiment(parsed_system, parsed_scenario, single_sbatch) + parsed_system.output_path = (output_dir or parsed_system.output_path).expanduser().resolve() + if prepare_output_dir(parsed_system.output_path) is None: + raise OSError(f"Cannot prepare output directory: {parsed_system.output_path}") + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + result_dir = Path( + tempfile.mkdtemp(prefix=f"{parsed_scenario.name}_{timestamp}_", dir=parsed_system.output_path) + ).resolve() + runner = create_experiment_runner( + parsed_system, parsed_scenario, dry_run=dry_run, single_sbatch=single_sbatch, result_dir=result_dir + ) + execute_experiment(runner, tests, enable_cache_without_check) + return runner.runner.experiment_output.snapshot() + + +def get_experiment(exp: str | Path) -> cloudai.models.output.Experiment: + """Read a saved experiment from its result directory or JSON file.""" + path = Path(exp).expanduser() + if path.is_dir(): + path /= "experiment.json" + return cloudai.models.output.Experiment.model_validate_json(path.read_text(encoding="utf-8")) + + +def list_experiments(system: Path) -> list[tuple[str, Path]]: + """Discover experiments beneath the configured results directory.""" + output_dir = Parser.parse_system(system.expanduser()).output_path.expanduser().resolve() + if not output_dir.exists(): + return [] + if not output_dir.is_dir(): + raise NotADirectoryError(output_dir) + return [(get_experiment(path).id, path.parent) for path in sorted(output_dir.glob("*/experiment.json"))] + + +def validate_scenario( + scenario: Path, + system: Path, + *, + tests_dir: Path | None = None, + hook_dir: Path | None = None, + single_sbatch: bool = False, +) -> tuple[bool, dict[str, str]]: + """Validate configuration and execution constraints without running workloads.""" + try: + parsed_system, _, parsed_scenario = load_scenario( + scenario.expanduser().resolve(), system.expanduser().resolve(), tests_dir=tests_dir, hook_dir=hook_dir + ) + validate_experiment(parsed_system, parsed_scenario, single_sbatch) + except SystemConfigParsingError as exc: + return False, {"system": str(exc.__cause__ or exc)} + except (TestConfigParsingError, TestScenarioParsingError, MissingTestError, ValueError, OSError) as exc: + return False, {"scenario": str(exc)} + return True, {} 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_api.py b/tests/test_api.py new file mode 100644 index 000000000..544158c2d --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,186 @@ +# 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 logging +import signal +from pathlib import Path +from unittest.mock import Mock + +import pydantic +import pytest +import toml + +import cloudai.api +import cloudai.core +import cloudai.handlers +import cloudai.metrics +from cloudai.systems.slurm import SingleSbatchRunner, SlurmRunner, SlurmSystem +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 configs(standalone_system: StandaloneSystem, tmp_path: Path) -> tuple[Path, Path]: + 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}, + } + ], + } + ) + ) + return scenario, system + + +@pytest.mark.parametrize("successful", [True, False]) +def test_experiment_api(configs: tuple[Path, Path], tmp_path: Path, monkeypatch: pytest.MonkeyPatch, successful: bool): + scenario, system = configs + 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(successful, "test outcome"), + ) + monkeypatch.setattr( + SleepTestDefinition, + "metric_observations", + lambda *_: [cloudai.metrics.MetricObservation(cloudai.metrics.LATENCY, 2.0, {})], + ) + handlers = logging.getLogger().handlers.copy() + sigint = signal.getsignal(signal.SIGINT) + assert cloudai.api.list_experiments(system) == [] + assert cloudai.api.validate_scenario(scenario, system) == (True, {}) + assert not (tmp_path / "output").exists() + experiment = cloudai.api.run_experiment(scenario, system) + assert logging.getLogger().handlers == handlers + assert signal.getsignal(signal.SIGINT) == sigint + assert experiment.status == ("completed" if successful else "failed") + assert experiment.start is not None and experiment.finish is not None + test = experiment.tests[0] + assert test.status == experiment.status + assert len(test.runs) == 1 + assert test.runs[0].status == experiment.status + expected_metrics = [{"name": "Latency", "value": 2.0, "unit": "us", "dimensions": []}] if successful else [] + assert [metric.model_dump() for metric in test.metrics] == expected_metrics + assert test.metrics == test.runs[0].metrics + result_dir = Path(experiment.path) + assert cloudai.api.get_experiment(result_dir / "experiment.json") == experiment + assert cloudai.api.list_experiments(system) == [(experiment.id, result_dir)] + again = cloudai.api.run_experiment(scenario, system) + assert again.id != experiment.id + assert len(cloudai.api.list_experiments(system)) == 2 + (result_dir / "experiment.json").write_text("{}") + with pytest.raises(pydantic.ValidationError): + cloudai.api.get_experiment(result_dir) + with pytest.raises(pydantic.ValidationError): + cloudai.api.list_experiments(system) + + +def test_validate_scenario(configs: tuple[Path, Path], tmp_path: Path): + scenario, system = configs + template = toml.loads(scenario.read_text())["Tests"][0] + template.pop("id") + (tmp_path / "sleep.toml").write_text(toml.dumps(template)) + relative = 'name = "relative"\n[[Tests]]\nid = "sleep"\npath = "sleep.toml"\n' + scenario.write_text(relative) + assert cloudai.api.validate_scenario(scenario, system) == (True, {}) + for invalid in ("name = [", 'name = "invalid"'): + scenario.write_text(invalid) + valid, errors = cloudai.api.validate_scenario(scenario, system) + assert not valid and errors["scenario"] + with pytest.raises(cloudai.core.TestScenarioParsingError): + cloudai.api.run_experiment(scenario, system) + scenario.write_text(relative) + original = system.read_text() + system.write_text(original.replace('scheduler = "standalone"', 'scheduler = "bad"')) + valid, errors = cloudai.api.validate_scenario(scenario, system) + assert not valid and "Unsupported system type" in errors["system"] + system.write_text(original.replace('name = "standalone"', "name = 42")) + valid, errors = cloudai.api.validate_scenario(scenario, system) + assert not valid and "name" in errors["system"] + system.write_text(original) + valid, errors = cloudai.api.validate_scenario(scenario, tmp_path / "missing.toml") + assert not valid and "missing.toml" in errors["system"] + valid, errors = cloudai.api.validate_scenario(scenario, system, single_sbatch=True) + assert not valid and "only supported for Slurm" in errors["scenario"] + assert not (tmp_path / "output").exists() + with pytest.raises(FileNotFoundError): + cloudai.api.get_experiment(tmp_path / "missing") + + +@pytest.mark.parametrize("failure", ["setup", "submission"]) +def test_failed_experiment_is_saved(configs: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, failure: str): + scenario, system = configs + if failure == "setup": + monkeypatch.setattr(StandaloneSystem, "update", Mock(side_effect=RuntimeError("setup failed"))) + else: + monkeypatch.setattr( + StandaloneRunner, + "_submit_test", + Mock( + side_effect=cloudai.core.JobSubmissionError( + test_name="sleep", command="sleep 0", stdout="", stderr="", message="submission failed" + ) + ), + ) + with pytest.raises((RuntimeError, cloudai.core.JobSubmissionError)): + cloudai.api.run_experiment(scenario, system) + experiments = cloudai.api.list_experiments(system) + assert len(experiments) == 1 + failed = cloudai.api.get_experiment(experiments[0][1]) + assert failed.status == "failed" + assert failed.finish is not None + assert failed.tests[0].runs == [] + + +def test_single_sbatch_selection_is_local(slurm_system: SlurmSystem, base_tr: cloudai.core.TestRun): + scenario = cloudai.core.TestScenario(name="selection", test_runs=[base_tr]) + registered = cloudai.core.Registry().runners_map["slurm"] + single = cloudai.handlers.create_experiment_runner(slurm_system, scenario, single_sbatch=True) + normal = cloudai.handlers.create_experiment_runner(slurm_system, scenario) + assert isinstance(single.runner, SingleSbatchRunner) + assert type(normal.runner) is SlurmRunner + assert cloudai.core.Registry().runners_map["slurm"] is registered diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 9f401ec73..ee49c111b 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, cast @@ -27,7 +26,7 @@ import cloudai.models.output from cloudai.cli.cli import verify_system_configs, verify_test_configs, verify_test_scenarios from cloudai.configurator import CloudAIGymEnv -from cloudai.configurator.env_params import EnvParamSpec, validate_domain_randomization_active +from cloudai.configurator.env_params import EnvParamSpec from cloudai.core import ( BaseAgent, BaseAgentConfig, @@ -44,6 +43,7 @@ from cloudai.handlers import ( handle_dse_job, prepare_installation, + validate_domain_randomization_active, ) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import CmdArgs, TestDefinition @@ -112,7 +112,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, "dry-run") == 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,7 +158,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, "dry-run") == 0 assert len(StubAgent.received_configs) == 1 recorded = StubAgent.received_configs[0] @@ -253,7 +253,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, "dry-run") == 0 reporter = StatusReporter( inner_runner.system, @@ -392,7 +392,7 @@ def test_handle_dse_job_invokes_agent_run( 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, "dry-run") == 0 assert CustomRunStubAgent.run_calls == 1 @@ -455,9 +455,9 @@ def run(agent: CustomRunStubAgent) -> int: monkeypatch.setattr(CustomRunStubAgent, "run", run) if failed_steps: with pytest.raises(RuntimeError, match="trial failed"): - handle_dse_job(runner, argparse.Namespace(mode="dry-run")) + handle_dse_job(runner, "dry-run") else: - assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 0 + assert handle_dse_job(runner, "dry-run") == 0 stored = cloudai.models.output.Experiment.model_validate_json( (runner.runner.scenario_root / "experiment.json").read_text() @@ -487,7 +487,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, "dry-run") == 1 assert CustomRunStubAgent.run_calls == 1 @@ -509,7 +509,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, "dry-run") == 1 assert CustomRunStubAgent.run_calls == 2 @@ -529,7 +529,7 @@ def test_handle_dse_job_propagates_agent_run_exception( 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, "dry-run") assert CustomRunStubAgent.run_calls == 1 @@ -547,7 +547,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, "dry-run") assert CustomRunStubAgent.run_calls == 1 @@ -565,7 +565,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, "run") failure_report = tmp_path / "dse_failure.txt" assert failure_report.exists()