diff --git a/doc/reporting.rst b/doc/reporting.rst index 87034e7ff..e55cb4d25 100644 --- a/doc/reporting.rst +++ b/doc/reporting.rst @@ -36,15 +36,26 @@ Unified experiment output CloudAI writes ``experiment.json`` in each scenario's results directory. -The file contains scenario details, test cases, status, timing, and result paths. Standalone execution also records each -run's process ID, status, iteration, timing, and workload metrics. +The file contains scenario details and test cases under ``tests``. Standalone and Slurm execution records appear in each +test case's ``runs`` list. Each record represents an iteration or DSE step and includes its number, process or Slurm job +ID, status, timing, result path, and workload metrics. -CloudAI updates the file when standalone runs start and finish, then finalizes it when scenario execution succeeds or -fails. Timestamps use UTC; durations use seconds. Unknown timestamps are ``null``. Dry runs also produce scenario and -test-case details without launching workloads. +Timestamps use UTC; durations use seconds. Unknown timestamps are ``null``. A final status of ``unknown`` means the outcome +could not be determined. Dry runs include scenario and test-case details without launching workloads. -Metrics come from ``TestDefinition.metric_observations()``, independently of reporter settings. Each file update replaces -the previous snapshot atomically. Metric extraction or write errors produce warnings without affecting execution. +When a test case executes once successfully, ``tests[].metrics`` contains that execution's metrics. For DSE, it contains +metrics from the successful step with the highest valid reward. The search space, selected step, and configuration appear +in ``tests[].dse``. For example: + +.. code-block:: json + + { + "space": {"extra_env_vars.NCCL_ALGO": ["Ring", "Tree"]}, + "best_step": 2, + "best_config": {"extra_env_vars.NCCL_ALGO": "Tree"} + } + +Here, step 2 using ``Tree`` was selected. Its metrics appear in the test case's ``metrics`` list. .. _general-flow: diff --git a/pyproject.toml b/pyproject.toml index 394878152..cb643ee67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -152,20 +152,26 @@ root_package = "cloudai" "cloudai.workloads", "cloudai.util", "cloudai.cli", + "cloudai.handlers", "cloudai.report_generator", ] [[tool.importlinter.contracts]] name = "Report generator is leaf dependency" type = "forbidden" - forbidden_modules = ["cloudai.workloads", "cloudai.cli"] - allow_indirect_imports = true # allow "from cloudai.core import ..." + forbidden_modules = ["cloudai.workloads", "cloudai.cli", "cloudai.handlers"] + allow_indirect_imports = true # allow "from cloudai.core import ..." source_modules = ["cloudai.report_generator"] [[tool.importlinter.contracts]] name = "Util modules are leaf dependencies" type = "forbidden" - forbidden_modules = ["cloudai.systems", "cloudai.workloads", "cloudai.cli"] + forbidden_modules = [ + "cloudai.systems", + "cloudai.workloads", + "cloudai.cli", + "cloudai.handlers", + ] allow_indirect_imports = true source_modules = ["cloudai.util"] diff --git a/src/cloudai/_core/base_runner.py b/src/cloudai/_core/base_runner.py index dbd8f70e7..681611c8b 100644 --- a/src/cloudai/_core/base_runner.py +++ b/src/cloudai/_core/base_runner.py @@ -113,7 +113,9 @@ def update_run_output(self, job: BaseJob, result: JobStatusResult | None = None) self.experiment_output.write() def finish_output(self, successful: bool) -> None: - status = "completed" if successful else "failed" + status: cloudai.models.output.Status = "completed" if successful else "failed" + if self.shutting_down and not any(test.status == "failed" for test in self.experiment_output.experiment.tests): + status = "cancelled" self.experiment_output.finish(status=status, finish=datetime.datetime.now(datetime.timezone.utc)) def shutdown(self): diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 0aa8acdad..b2c920a95 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -17,17 +17,35 @@ import argparse import logging import logging.config +import signal +from contextlib import contextmanager from pathlib import Path +from typing import Callable, List +from unittest.mock import Mock import click +import toml +import yaml -from .handlers import ( - handle_dry_run_and_run, - handle_generate_report, - handle_install_and_uninstall, - handle_list_registered_items, - handle_verify_all_configs, +import cloudai.handlers +from cloudai.configurator.env_params import validate_domain_randomization_active +from cloudai.core import ( + InstallStatusResult, + MissingTestError, + Parser, + Registry, + Runner, + System, + TestDefinition, + TestParser, + TestScenario, + TestScenarioParsingError, ) +from cloudai.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: @@ -83,6 +101,471 @@ def setup_logging(log_file: str, log_level: str) -> None: logging.config.dictConfig(LOGGING_CONFIG) +def _log_installation_dirs(prefix: str, system: System) -> None: + logging.info(f"{prefix} '{system.install_path.absolute()}'. HF cache is {system.hf_home_path.absolute()}.") + + +def handle_install_and_uninstall(args: argparse.Namespace) -> int: + """ + Manage the installation or uninstallation process for CloudAI. + + Based on user-specified mode, utilizing the Installer class. + + Args: + args (argparse.Namespace): The parsed command-line arguments. + """ + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) + system, tests, scenario = parser.parse(args.tests_dir, args.test_scenario) + + system.update() + logging.info(f"System Name: {system.name}") + logging.info(f"Scheduler: {system.scheduler}") + + installables, installer = cloudai.handlers.prepare_installation(system, tests, scenario) + + rc = 0 + if args.mode == "install": + all_installed = installer.is_installed(installables) + if all_installed: + _log_installation_dirs("CloudAI is already installed into", system) + else: + logging.info("Not all components are ready") + result = installer.install(installables) + if result.success: + _log_installation_dirs("CloudAI is successfully installed into", system) + else: + logging.error(result.message) + rc = 1 + + elif args.mode == "uninstall": + logging.info("Uninstalling test templates.") + result = installer.uninstall(installables) + if result.success: + logging.info("Uninstallation successful.") + else: + logging.error(result.message) + rc = 1 + + return rc + + +def register_signal_handlers(signal_handler: Callable) -> None: + """Register signal handlers for handling termination-related signals.""" + signals = [ + signal.SIGINT, + signal.SIGTERM, + signal.SIGHUP, + signal.SIGQUIT, + ] + for sig in signals: + signal.signal(sig, signal_handler) + + +def handle_generate_report(args: argparse.Namespace) -> int: + """ + Generate a report based on the existing configuration and test results. + + Args: + args (argparse.Namespace): The parsed command-line arguments. + """ + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) + system, _, test_scenario = parser.parse(args.tests_dir, args.test_scenario) + assert test_scenario is not None + + cloudai.handlers.generate_reports(system, test_scenario, args.result_dir) + + logging.debug("Report generation completed.") + + return 0 + + +def expand_file_list(root: Path, glob: str = "*.toml") -> tuple[int, List[Path]]: + if not root.exists(): + logging.error(f"{root} does not exist.") + return (1, []) + + test_tomls = [root] + if root.is_dir(): + test_tomls = list(root.glob(glob)) + if not test_tomls: + logging.error(f"No TOMLs found in {root}") + return (1, []) + + return (0, test_tomls) + + +@contextmanager +def _ensure_kube_config_exists(system_toml_path: Path, content: str): + try: + config_dict = toml.loads(content) + except toml.TomlDecodeError as e: + logging.error(format_toml_decode_error(system_toml_path, e, "system config")) + raise + + kube_config_path_str = config_dict.get("kube_config_path") + kube_config_path = Path(kube_config_path_str) if kube_config_path_str else Path.home() / ".kube" / "config" + + created_file = False + created_dir = False + + if not kube_config_path.exists(): + logging.warning(f"Kube config file '{kube_config_path}' not found. Creating a dummy one.") + if not kube_config_path.parent.exists(): + kube_config_path.parent.mkdir(parents=True, exist_ok=True) + created_dir = True + + dummy_config = { + "apiVersion": "v1", + "kind": "Config", + "preferences": {}, + "clusters": [{"name": "dummy-cluster", "cluster": {"server": "https://dummy-server"}}], + "users": [{"name": "dummy-user", "user": {"token": "dummy-token"}}], + "contexts": [{"name": "dummy-context", "context": {"cluster": "dummy-cluster", "user": "dummy-user"}}], + "current-context": "dummy-context", + } + kube_config_path.write_text(yaml.dump(dummy_config)) + created_file = True + else: + logging.debug(f"Kube config '{kube_config_path}' already exists. Skipping creation.") + + try: + yield kube_config_path + finally: + if created_file: + try: + kube_config_path.unlink() + logging.debug(f"Deleted temporary kube config: {kube_config_path}") + except Exception as e: + logging.warning(f"Failed to remove temporary kube config '{kube_config_path}': {e}") + if created_dir: + try: + kube_config_path.parent.rmdir() + logging.debug(f"Deleted kube config directory: {kube_config_path.parent}") + except OSError: + pass + + +def verify_system_configs(system_tomls: List[Path]) -> int: + nfailed = 0 + + for system_toml in system_tomls: + logging.debug(f"Verifying System: {system_toml}...") + content = system_toml.read_text() + + if 'scheduler = "kubernetes"' in content: + try: + with _ensure_kube_config_exists(system_toml, content): + Parser.parse_system(system_toml) + except Exception as e: + logging.error(f"Failed to verify system config {system_toml}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + else: + try: + Parser.parse_system(system_toml) + except Exception as e: + logging.error(f"Failed to verify system config {system_toml}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + + if nfailed: + logging.error(f"{nfailed} out of {len(system_tomls)} system configurations have issues.") + else: + logging.info(f"Checked systems: {len(system_tomls)}, all passed") + + return nfailed + + +def verify_test_configs(test_tomls: List[Path]) -> int: + nfailed = 0 + tp = TestParser([], None) # type: ignore + for test_toml in test_tomls: + logging.debug(f"Verifying Test: {test_toml}...") + try: + with test_toml.open() as fh: + tp.current_file = test_toml + tp.load_test_definition(load_test_toml_file(fh, test_toml)) + except Exception as e: + logging.error(f"Failed to verify Test: {test_toml}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + + if nfailed: + logging.error(f"{nfailed} out of {len(test_tomls)} test configurations have issues.") + else: + logging.info(f"Checked tests: {len(test_tomls)}, all passed") + + return nfailed + + +def verify_test_scenarios( + scenario_tomls: List[Path], test_tomls: list[Path], hook_tomls: List[Path], hook_test_tomls: list[Path] +) -> int: + system = Mock(spec=System, sol={}) + nfailed = 0 + for scenario_file in scenario_tomls: + logging.debug(f"Verifying Test Scenario: {scenario_file}...") + try: + tests = Parser.parse_tests(test_tomls, system) + hook_tests = Parser.parse_tests(hook_test_tomls, system) + hooks = Parser.parse_hooks(hook_tomls, system, {t.name: t for t in hook_tests}) + scenario = Parser.parse_test_scenario(scenario_file, system, {t.name: t for t in tests}, hooks) + validate_domain_randomization_active(scenario) + except Exception as e: + logging.error(f"Failed to verify Test Scenario: {scenario_file}: {e}") + logging.debug("", exc_info=True) + nfailed += 1 + + if nfailed: + logging.error(f"{nfailed} out of {len(scenario_tomls)} test scenarios have issues.") + else: + logging.info(f"Checked scenarios: {len(scenario_tomls)}, all passed") + + return nfailed + + +def handle_verify_all_configs(args: argparse.Namespace) -> int: + root: Path = args.configs_dir + err, tomls = expand_file_list(root, glob="**/*.toml") + if err: + return err + + err, hook_tomls = expand_file_list(HOOK_ROOT, glob="**/*.toml") + tomls += hook_tomls + logging.info(f"Found {len(hook_tomls)} hook TOMLs (always verified)") + + files = load_tomls_by_type(tomls) + + test_tomls = files["test"] + if args.tests_dir: + test_tomls = list(args.tests_dir.glob("*.toml")) + elif files["scenario"]: + logging.warning( + "Test configuration directory not provided, using all found test TOMLs in the specified directory." + ) + + nfailed = 0 + total_checked = 0 + + if files["system"]: + nfailed += verify_system_configs(files["system"]) + total_checked += len(files["system"]) + if test_tomls: + nfailed += verify_test_configs(test_tomls) + total_checked += len(test_tomls) + if files["scenario"]: + nfailed += verify_test_scenarios(files["scenario"], test_tomls, files["hook"], files["hook_test"]) + total_checked += len(files["scenario"]) + if files["unknown"]: + for unknown_file in files["unknown"]: + logging.error( + f"Unknown configuration file '{unknown_file}': could not classify as system, test, scenario, or hook." + ) + nfailed += len(files["unknown"]) + total_checked += len(files["unknown"]) + + if nfailed: + logging.error(f"{nfailed} out of {total_checked} configuration files have issues.") + else: + logging.info(f"Checked {total_checked} configuration files, all passed") + + return nfailed + + +def load_tomls_by_type(tomls: List[Path]) -> dict[str, List[Path]]: + files: dict[str, List[Path]] = { + "system": [], + "test": [], + "scenario": [], + "hook_test": [], + "hook": [], + "unknown": [], + } + for toml_file in tomls: + content = toml_file.read_text() + + is_in_hook_root = False + try: + toml_file.relative_to(HOOK_ROOT) + is_in_hook_root = True + except ValueError: + pass + + if is_in_hook_root: + if "test" in toml_file.parts: + files["hook_test"].append(toml_file) + else: + files["hook"].append(toml_file) + continue + + try: + toml_content = toml.loads(content) + except toml.TomlDecodeError: + files["unknown"].append(toml_file) + continue + + if not isinstance(toml_content, dict): + files["unknown"].append(toml_file) + continue + + if "scheduler" in toml_content: + files["system"].append(toml_file) + elif "test_template_name =" in content and "[[Tests]]" not in content: + files["test"].append(toml_file) + elif "[[Tests]]" in content: + files["scenario"].append(toml_file) + else: + files["unknown"].append(toml_file) + + return files + + +def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: C901 + registry = Registry() + if item_type.lower() == "reports": + print("Available scenario reports:") + for idx, (name, report) in enumerate(sorted(registry.scenario_reports.items()), start=1): + string = f'{idx}. "{name}" {report.__name__}' + if verbose: + string += f" (config={registry.report_configs[name].model_dump_json(indent=None)})" + print(string) + elif item_type.lower() == "agents": + print("Available agents:") + for idx, name in enumerate(registry.agent_names(), start=1): + agent = registry.get_agent(name) + string = f'{idx}. "{name}" class={agent.__name__}' + if verbose: + string += f"{agent.__doc__}" + print(string) + elif item_type.lower() == "reward-functions": + print("Available reward functions:") + for idx, name in enumerate(registry.reward_function_names(), start=1): + reward_function = registry.get_reward_function(name) + callable_name = getattr(reward_function, "__name__", type(reward_function).__name__) + string = f'{idx}. "{name}" function={callable_name}' + if verbose: + documentation = getattr(reward_function, "__doc__", None) + if documentation: + string += f" {documentation}" + print(string) + + return 0 + + +def _setup_system_and_scenario( + args: argparse.Namespace, +) -> tuple[System, TestScenario, list[TestDefinition]] | None: + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) + try: + system, tests, test_scenario = parser.parse(args.tests_dir, args.test_scenario) + except MissingTestError as e: + logging.error(e.message) + return None + + assert test_scenario is not None + + if args.output_dir: + system.output_path = args.output_dir.absolute() + + if not prepare_output_dir(system.output_path): + return None + + if args.mode == "dry-run": + system.monitor_interval = 1 + system.update() + + return system, test_scenario, tests + + +def _handle_single_sbatch(args: argparse.Namespace, system: System) -> bool: + if not args.single_sbatch: + return True + + if not isinstance(system, SlurmSystem): + logging.error("Single sbatch is only supported for Slurm systems.") + return False + + Registry().update_runner("slurm", SingleSbatchRunner) + return True + + +def _check_installation( + args: argparse.Namespace, system: System, tests: list[TestDefinition], test_scenario: TestScenario +) -> InstallStatusResult: + logging.info("Checking if workloads components are installed.") + installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) + + if args.enable_cache_without_check: + result = installer.mark_as_installed(installables) + else: + result = installer.is_installed(installables) + + return result + + +def handle_dry_run_and_run(args: argparse.Namespace) -> int: + setup_result = _setup_system_and_scenario(args) + if setup_result is None: + return 1 + system, test_scenario, tests = setup_result + + try: + validate_domain_randomization_active(test_scenario) + except TestScenarioParsingError as e: + logging.error(str(e)) + return 1 + + if not _handle_single_sbatch(args, system): + return 1 + + logging.info(f"System Name: {system.name}") + logging.info(f"Scheduler: {system.scheduler}") + logging.info(f"Test Scenario Name: {test_scenario.name}") + + result = _check_installation(args, system, tests, test_scenario) + if args.mode == "run" and not result.success: + logging.info("Not all workloads components are installed. Installing...") + installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) + + result = installer.install(installables) + if result.success: + _log_installation_dirs("CloudAI is successfully installed into", system) + else: + logging.error("Failed to install workloads components.") + logging.error(result.message) + return 1 + elif args.mode == "dry-run": + # simulate installation for dry-run + installables, installer = cloudai.handlers.prepare_installation(system, tests, test_scenario) + result = installer.mark_as_installed(installables) + if not result.success: + logging.warning("Failed to mark workloads components as installed for dry-run.") + + logging.info(test_scenario.pretty_print()) + + runner = Runner(args.mode, system, test_scenario) + register_signal_handlers(runner.cancel_on_signal) + logging.info(f"Scenario results will be stored at: {runner.runner.scenario_root}") + + successful = False + try: + runner.runner.experiment_output.write() + has_dse = any(tr.is_dse_job for tr in test_scenario.test_runs) + if args.single_sbatch or not has_dse: # in this mode cases are unrolled using grid search + successful = cloudai.handlers.handle_non_dse_job(runner, args) + return 0 + + if all(tr.is_dse_job for tr in test_scenario.test_runs): + result = cloudai.handlers.handle_dse_job(runner, args) + successful = result == 0 + return result + + logging.error("Mixing DSE and non-DSE jobs is not allowed.") + return 1 + finally: + runner.runner.finish_output(successful) + + def common_options(f): f = click.option( "--system-config", @@ -271,9 +754,9 @@ def verify_configs(configs_dir: Path, tests_dir: Path): raise click.exceptions.Exit(1 if handle_verify_all_configs(args) else 0) -@main.command() +@main.command(name="list") @click.argument("type", type=click.Choice(["reports", "agents", "reward-functions"], case_sensitive=False)) @click.option("-v", "--verbose", is_flag=True, default=False, help="Verbose output.") -def list(type: str, verbose: bool): +def list_registered(type: str, verbose: bool): """List available in Registry items.""" handle_list_registered_items(type, verbose) diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 313cbbed2..7e27d172d 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -14,650 +14,51 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse -import copy -import logging -import signal -import traceback -from contextlib import contextmanager -from pathlib import Path -from typing import Callable, List, Optional -from unittest.mock import Mock +import functools +import typing +import warnings -import toml -import yaml +import cloudai.cli.cli +import cloudai.configurator.env_params +import cloudai.handlers +import cloudai.test_parser +import cloudai.toml_utils +import cloudai.util -from cloudai.configurator.env_params import validate_domain_randomization_active -from cloudai.core import ( - BaseInstaller, - CloudAIGymEnv, - Installable, - InstallStatusResult, - MissingTestError, - Parser, - Registry, - Runner, - System, - TestParser, - TestScenario, - TestScenarioParsingError, -) -from cloudai.models.scenario import ReportConfig -from cloudai.models.workload import TestDefinition -from cloudai.parser import HOOK_ROOT -from cloudai.systems.slurm import SingleSbatchRunner, SlurmSystem -from cloudai.test_parser import load_test_toml_file -from cloudai.toml_utils import format_toml_decode_error -from cloudai.util import prepare_output_dir +_P = typing.ParamSpec("_P") +_R = typing.TypeVar("_R") -def _log_installation_dirs(prefix: str, system: System) -> None: - logging.info(f"{prefix} '{system.install_path.absolute()}'. HF cache is {system.hf_home_path.absolute()}.") - - -def handle_install_and_uninstall(args: argparse.Namespace) -> int: - """ - Manage the installation or uninstallation process for CloudAI. - - Based on user-specified mode, utilizing the Installer class. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - """ - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - system, tests, scenario = parser.parse(args.tests_dir, args.test_scenario) - - system.update() - logging.info(f"System Name: {system.name}") - logging.info(f"Scheduler: {system.scheduler}") - - installables, installer = prepare_installation(system, tests, scenario) - - rc = 0 - if args.mode == "install": - all_installed = installer.is_installed(installables) - if all_installed: - _log_installation_dirs("CloudAI is already installed into", system) - else: - logging.info("Not all components are ready") - result = installer.install(installables) - if result.success: - _log_installation_dirs("CloudAI is successfully installed into", system) - else: - logging.error(result.message) - rc = 1 - - elif args.mode == "uninstall": - logging.info("Uninstalling test templates.") - result = installer.uninstall(installables) - if result.success: - logging.info("Uninstallation successful.") - else: - logging.error(result.message) - rc = 1 - - return rc - - -def prepare_installation( - system: System, tests: list[TestDefinition], scenario: Optional[TestScenario] -) -> tuple[list[Installable], BaseInstaller]: - installables: list[Installable] = [] - if scenario: - installables.extend(_scenario_installables(scenario)) - else: - for test in tests: - logging.debug(f"{test.name} has {len(test.installables)} installables.") - installables.extend(test.installables) - - registry = Registry() - installer_class = registry.installers_map.get(system.scheduler, BaseInstaller) - if installer_class is None: - raise NotImplementedError(f"No installer available for scheduler: {system.scheduler}") - installer = installer_class(system) - - return installables, installer - - -def _scenario_installables(scenario: TestScenario) -> list[Installable]: - installables: list[Installable] = [] - for test_run in scenario.test_runs: - logging.debug(f"{test_run.test.name} has {len(test_run.test.installables)} installables.") - installables.extend(test_run.test.installables) - for hook in (test_run.pre_test, test_run.post_test): - if hook is not None: - installables.extend(_scenario_installables(hook)) - return installables - - -def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: - registry = Registry() - - original_test_runs = copy.deepcopy(runner.runner.test_scenario.test_runs) - - has_dependencies = any(tr.dependencies for tr in runner.runner.test_scenario.test_runs) - if has_dependencies: - logging.error( - "Dependencies are not supported for DSE jobs, all cases run consecutively. " - "Please remove dependencies and re-run." - ) - return 1 - - err = 0 - # Capture an unexpected error so reports still generate, then re-raise below. - run_error: Exception | None = None - try: - for tr in runner.runner.test_scenario.test_runs: - test_run = copy.deepcopy(tr) - - agent_type = test_run.test.agent - if not registry.has_agent(agent_type): - logging.error( - f"No agent available for type: {agent_type}. Please make sure {agent_type} " - f"is a valid agent type. Available agents: {registry.agent_names()}" - ) - err = 1 - continue - agent_class = registry.get_agent(agent_type) - - agent_config_data = test_run.test.agent_config or {} - agent_config = agent_class.get_config_class()(**agent_config_data) - env = CloudAIGymEnv( - test_run=test_run, - runner=runner.runner, - rewards=agent_config.rewards, - ) - if agent_config.start_action == "first": - logging.info(f"Using deterministic first sweep for the chosen agent: {env.first_sweep}.") - - agent = agent_class(env, agent_config) - logging.debug(f"Created agent {agent.__class__.__name__}.") - - err |= agent.run() - except Exception as exc: - run_error = exc - logging.exception("DSE job aborted by an unexpected error; generating reports before failing.") - - if args.mode == "run": - runner.runner.test_scenario.test_runs = original_test_runs - generate_reports( - runner.runner.system, - runner.runner.test_scenario, - runner.runner.scenario_root, - error=run_error, - ) - - if run_error is not None: - raise run_error.with_traceback(run_error.__traceback__) - - logging.info("All jobs are complete.") - return err - - -def _record_run_failure(result_dir: Path, error: BaseException) -> None: - """Persist an aborting error into the results dir so the failure is documented with the reports.""" - failure_path = result_dir / "dse_failure.txt" - tb = "".join(traceback.format_exception(type(error), error, error.__traceback__)) - try: - result_dir.mkdir(parents=True, exist_ok=True) - failure_path.write_text(f"DSE job aborted by an unexpected {type(error).__name__}: {error}\n\n{tb}") - logging.info(f"Documented the aborting error in {failure_path}") - except OSError: - logging.exception(f"Failed to write failure report to {failure_path}") - - -def generate_reports( - system: System, - test_scenario: TestScenario, - result_dir: Path, - error: BaseException | None = None, -) -> None: - registry = Registry() - - if error is not None: - _record_run_failure(result_dir, error) - - for name, reporter_class in registry.ordered_scenario_reports(): - logging.debug(f"Generating report '{name}' ({reporter_class.__name__})") - - cfg = registry.report_configs.get(name, ReportConfig(enable=False)) - if scenario_cfg := test_scenario.reports.get(name): - cfg = scenario_cfg - elif isinstance(system, SlurmSystem) and system.reports and name in system.reports: - cfg = system.reports[name] - logging.debug(f"Report '{name}' config is: {cfg.model_dump_json(indent=None)}") - - if not cfg.enable: - logging.debug(f"Skipping report {name} because it is disabled.") - continue - - try: - reporter = reporter_class(system, test_scenario, result_dir, cfg) - reporter.generate() - except Exception as e: - logging.warning(f"Error generating report '{name}', see debug log for details") - logging.debug(e, exc_info=True) - - -def handle_non_dse_job(runner: Runner, args: argparse.Namespace) -> 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) - - -def _setup_system_and_scenario( - args: argparse.Namespace, -) -> tuple[System, TestScenario, list[TestDefinition]] | None: - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - try: - system, tests, test_scenario = parser.parse(args.tests_dir, args.test_scenario) - except MissingTestError as e: - logging.error(e.message) - return None - - assert test_scenario is not None - - if args.output_dir: - system.output_path = args.output_dir.absolute() - - if not prepare_output_dir(system.output_path): - return None - - if args.mode == "dry-run": - system.monitor_interval = 1 - system.update() - - return system, test_scenario, tests - - -def _handle_single_sbatch(args: argparse.Namespace, system: System) -> bool: - if not args.single_sbatch: - return True - - if not isinstance(system, SlurmSystem): - logging.error("Single sbatch is only supported for Slurm systems.") - return False - - Registry().update_runner("slurm", SingleSbatchRunner) - return True - - -def _check_installation( - args: argparse.Namespace, system: System, tests: list[TestDefinition], test_scenario: TestScenario -) -> InstallStatusResult: - logging.info("Checking if workloads components are installed.") - installables, installer = prepare_installation(system, tests, test_scenario) - - if args.enable_cache_without_check: - result = installer.mark_as_installed(installables) - else: - result = installer.is_installed(installables) - - return result - - -def handle_dry_run_and_run(args: argparse.Namespace) -> int: - setup_result = _setup_system_and_scenario(args) - if setup_result is None: - return 1 - system, test_scenario, tests = setup_result - - try: - validate_domain_randomization_active(test_scenario) - except TestScenarioParsingError as e: - logging.error(str(e)) - return 1 - - if not _handle_single_sbatch(args, system): - return 1 - - logging.info(f"System Name: {system.name}") - logging.info(f"Scheduler: {system.scheduler}") - logging.info(f"Test Scenario Name: {test_scenario.name}") - - result = _check_installation(args, system, tests, test_scenario) - if args.mode == "run" and not result.success: - logging.info("Not all workloads components are installed. Installing...") - installables, installer = prepare_installation(system, tests, test_scenario) - - result = installer.install(installables) - if result.success: - _log_installation_dirs("CloudAI is successfully installed into", system) - else: - logging.error("Failed to install workloads components.") - logging.error(result.message) - return 1 - elif args.mode == "dry-run": - # simulate installation for dry-run - installables, installer = prepare_installation(system, tests, test_scenario) - result = installer.mark_as_installed(installables) - if not result.success: - logging.warning("Failed to mark workloads components as installed for dry-run.") - - logging.info(test_scenario.pretty_print()) - - runner = Runner(args.mode, system, test_scenario) - register_signal_handlers(runner.cancel_on_signal) - logging.info(f"Scenario results will be stored at: {runner.runner.scenario_root}") - - successful = False - try: - runner.runner.experiment_output.write() - has_dse = any(tr.is_dse_job for tr in test_scenario.test_runs) - if args.single_sbatch or not has_dse: # in this mode cases are unrolled using grid search - successful = handle_non_dse_job(runner, args) - return 0 - - if all(tr.is_dse_job for tr in test_scenario.test_runs): - result = handle_dse_job(runner, args) - successful = result == 0 - return result - - logging.error("Mixing DSE and non-DSE jobs is not allowed.") - return 1 - finally: - runner.runner.finish_output(successful) - - -def handle_generate_report(args: argparse.Namespace) -> int: - """ - Generate a report based on the existing configuration and test results. - - Args: - args (argparse.Namespace): The parsed command-line arguments. - """ - parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) - system, _, test_scenario = parser.parse(args.tests_dir, args.test_scenario) - assert test_scenario is not None - - generate_reports(system, test_scenario, args.result_dir) - - logging.debug("Report generation completed.") - - return 0 - - -def expand_file_list(root: Path, glob: str = "*.toml") -> tuple[int, List[Path]]: - if not root.exists(): - logging.error(f"{root} does not exist.") - return (1, []) - - test_tomls = [root] - if root.is_dir(): - test_tomls = list(root.glob(glob)) - if not test_tomls: - logging.error(f"No TOMLs found in {root}") - return (1, []) - - return (0, test_tomls) - - -@contextmanager -def _ensure_kube_config_exists(system_toml_path: Path, content: str): - try: - config_dict = toml.loads(content) - except toml.TomlDecodeError as e: - logging.error(format_toml_decode_error(system_toml_path, e, "system config")) - raise - - kube_config_path_str = config_dict.get("kube_config_path") - kube_config_path = Path(kube_config_path_str) if kube_config_path_str else Path.home() / ".kube" / "config" - - created_file = False - created_dir = False - - if not kube_config_path.exists(): - logging.warning(f"Kube config file '{kube_config_path}' not found. Creating a dummy one.") - if not kube_config_path.parent.exists(): - kube_config_path.parent.mkdir(parents=True, exist_ok=True) - created_dir = True - - dummy_config = { - "apiVersion": "v1", - "kind": "Config", - "preferences": {}, - "clusters": [{"name": "dummy-cluster", "cluster": {"server": "https://dummy-server"}}], - "users": [{"name": "dummy-user", "user": {"token": "dummy-token"}}], - "contexts": [{"name": "dummy-context", "context": {"cluster": "dummy-cluster", "user": "dummy-user"}}], - "current-context": "dummy-context", - } - kube_config_path.write_text(yaml.dump(dummy_config)) - created_file = True - else: - logging.debug(f"Kube config '{kube_config_path}' already exists. Skipping creation.") - - try: - yield kube_config_path - finally: - if created_file: - try: - kube_config_path.unlink() - logging.debug(f"Deleted temporary kube config: {kube_config_path}") - except Exception as e: - logging.warning(f"Failed to remove temporary kube config '{kube_config_path}': {e}") - if created_dir: - try: - kube_config_path.parent.rmdir() - logging.debug(f"Deleted kube config directory: {kube_config_path.parent}") - except OSError: - pass - - -def verify_system_configs(system_tomls: List[Path]) -> int: - nfailed = 0 - - for system_toml in system_tomls: - logging.debug(f"Verifying System: {system_toml}...") - content = system_toml.read_text() - - if 'scheduler = "kubernetes"' in content: - try: - with _ensure_kube_config_exists(system_toml, content): - Parser.parse_system(system_toml) - except Exception as e: - logging.error(f"Failed to verify system config {system_toml}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - else: - try: - Parser.parse_system(system_toml) - except Exception as e: - logging.error(f"Failed to verify system config {system_toml}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - - if nfailed: - logging.error(f"{nfailed} out of {len(system_tomls)} system configurations have issues.") - else: - logging.info(f"Checked systems: {len(system_tomls)}, all passed") - - return nfailed - - -def verify_test_configs(test_tomls: List[Path]) -> int: - nfailed = 0 - tp = TestParser([], None) # type: ignore - for test_toml in test_tomls: - logging.debug(f"Verifying Test: {test_toml}...") - try: - with test_toml.open() as fh: - tp.current_file = test_toml - tp.load_test_definition(load_test_toml_file(fh, test_toml)) - except Exception as e: - logging.error(f"Failed to verify Test: {test_toml}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - - if nfailed: - logging.error(f"{nfailed} out of {len(test_tomls)} test configurations have issues.") - else: - logging.info(f"Checked tests: {len(test_tomls)}, all passed") - - return nfailed - - -def verify_test_scenarios( - scenario_tomls: List[Path], test_tomls: list[Path], hook_tomls: List[Path], hook_test_tomls: list[Path] -) -> int: - system = Mock(spec=System, sol={}) - nfailed = 0 - for scenario_file in scenario_tomls: - logging.debug(f"Verifying Test Scenario: {scenario_file}...") - try: - tests = Parser.parse_tests(test_tomls, system) - hook_tests = Parser.parse_tests(hook_test_tomls, system) - hooks = Parser.parse_hooks(hook_tomls, system, {t.name: t for t in hook_tests}) - scenario = Parser.parse_test_scenario(scenario_file, system, {t.name: t for t in tests}, hooks) - validate_domain_randomization_active(scenario) - except Exception as e: - logging.error(f"Failed to verify Test Scenario: {scenario_file}: {e}") - logging.debug("", exc_info=True) - nfailed += 1 - - if nfailed: - logging.error(f"{nfailed} out of {len(scenario_tomls)} test scenarios have issues.") - else: - logging.info(f"Checked scenarios: {len(scenario_tomls)}, all passed") - - return nfailed - - -def handle_verify_all_configs(args: argparse.Namespace) -> int: - root: Path = args.configs_dir - err, tomls = expand_file_list(root, glob="**/*.toml") - if err: - return err - - err, hook_tomls = expand_file_list(HOOK_ROOT, glob="**/*.toml") - tomls += hook_tomls - logging.info(f"Found {len(hook_tomls)} hook TOMLs (always verified)") - - files = load_tomls_by_type(tomls) - - test_tomls = files["test"] - if args.tests_dir: - test_tomls = list(args.tests_dir.glob("*.toml")) - elif files["scenario"]: - logging.warning( - "Test configuration directory not provided, using all found test TOMLs in the specified directory." +def _deprecated(function: typing.Callable[_P, _R]) -> typing.Callable[_P, _R]: + @functools.wraps(function) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + warnings.warn( + f"cloudai.cli.handlers.{function.__name__} is deprecated; do not call CLI handlers directly.", + DeprecationWarning, + stacklevel=2, ) - - nfailed = 0 - total_checked = 0 - - if files["system"]: - nfailed += verify_system_configs(files["system"]) - total_checked += len(files["system"]) - if test_tomls: - nfailed += verify_test_configs(test_tomls) - total_checked += len(test_tomls) - if files["scenario"]: - nfailed += verify_test_scenarios(files["scenario"], test_tomls, files["hook"], files["hook_test"]) - total_checked += len(files["scenario"]) - if files["unknown"]: - for unknown_file in files["unknown"]: - logging.error( - f"Unknown configuration file '{unknown_file}': could not classify as system, test, scenario, or hook." - ) - nfailed += len(files["unknown"]) - total_checked += len(files["unknown"]) - - if nfailed: - logging.error(f"{nfailed} out of {total_checked} configuration files have issues.") - else: - logging.info(f"Checked {total_checked} configuration files, all passed") - - return nfailed - - -def load_tomls_by_type(tomls: List[Path]) -> dict[str, List[Path]]: - files: dict[str, List[Path]] = { - "system": [], - "test": [], - "scenario": [], - "hook_test": [], - "hook": [], - "unknown": [], - } - for toml_file in tomls: - content = toml_file.read_text() - - is_in_hook_root = False - try: - toml_file.relative_to(HOOK_ROOT) - is_in_hook_root = True - except ValueError: - pass - - if is_in_hook_root: - if "test" in toml_file.parts: - files["hook_test"].append(toml_file) - else: - files["hook"].append(toml_file) - continue - - try: - toml_content = toml.loads(content) - except toml.TomlDecodeError: - files["unknown"].append(toml_file) - continue - - if not isinstance(toml_content, dict): - files["unknown"].append(toml_file) - continue - - if "scheduler" in toml_content: - files["system"].append(toml_file) - elif "test_template_name =" in content and "[[Tests]]" not in content: - files["test"].append(toml_file) - elif "[[Tests]]" in content: - files["scenario"].append(toml_file) - else: - files["unknown"].append(toml_file) - - return files - - -def handle_list_registered_items(item_type: str, verbose: bool) -> int: # noqa: C901 - registry = Registry() - if item_type.lower() == "reports": - print("Available scenario reports:") - for idx, (name, report) in enumerate(sorted(registry.scenario_reports.items()), start=1): - string = f'{idx}. "{name}" {report.__name__}' - if verbose: - string += f" (config={registry.report_configs[name].model_dump_json(indent=None)})" - print(string) - elif item_type.lower() == "agents": - print("Available agents:") - for idx, name in enumerate(registry.agent_names(), start=1): - agent = registry.get_agent(name) - string = f'{idx}. "{name}" class={agent.__name__}' - if verbose: - string += f"{agent.__doc__}" - print(string) - elif item_type.lower() == "reward-functions": - print("Available reward functions:") - for idx, name in enumerate(registry.reward_function_names(), start=1): - reward_function = registry.get_reward_function(name) - callable_name = getattr(reward_function, "__name__", type(reward_function).__name__) - string = f'{idx}. "{name}" function={callable_name}' - if verbose: - documentation = getattr(reward_function, "__doc__", None) - if documentation: - string += f" {documentation}" - print(string) - - return 0 + return function(*args, **kwargs) + + return wrapper + + +handle_install_and_uninstall = _deprecated(cloudai.cli.cli.handle_install_and_uninstall) +prepare_installation = _deprecated(cloudai.handlers.prepare_installation) +generate_reports = _deprecated(cloudai.handlers.generate_reports) +register_signal_handlers = _deprecated(cloudai.cli.cli.register_signal_handlers) +handle_dry_run_and_run = _deprecated(cloudai.cli.cli.handle_dry_run_and_run) +handle_generate_report = _deprecated(cloudai.cli.cli.handle_generate_report) +expand_file_list = _deprecated(cloudai.cli.cli.expand_file_list) +verify_system_configs = _deprecated(cloudai.cli.cli.verify_system_configs) +verify_test_configs = _deprecated(cloudai.cli.cli.verify_test_configs) +verify_test_scenarios = _deprecated(cloudai.cli.cli.verify_test_scenarios) +handle_verify_all_configs = _deprecated(cloudai.cli.cli.handle_verify_all_configs) +load_tomls_by_type = _deprecated(cloudai.cli.cli.load_tomls_by_type) +handle_list_registered_items = _deprecated(cloudai.cli.cli.handle_list_registered_items) +validate_domain_randomization_active = _deprecated(cloudai.configurator.env_params.validate_domain_randomization_active) +load_test_toml_file = _deprecated(cloudai.test_parser.load_test_toml_file) +format_toml_decode_error = _deprecated(cloudai.toml_utils.format_toml_decode_error) +prepare_output_dir = _deprecated(cloudai.util.prepare_output_dir) + +handle_dse_job = _deprecated(cloudai.handlers.handle_dse_job) +handle_non_dse_job = _deprecated(cloudai.handlers.handle_non_dse_job) diff --git a/src/cloudai/configurator/cloudai_gym.py b/src/cloudai/configurator/cloudai_gym.py index cdafd5ce5..a472da9b5 100644 --- a/src/cloudai/configurator/cloudai_gym.py +++ b/src/cloudai/configurator/cloudai_gym.py @@ -16,6 +16,7 @@ import copy import logging +import math from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast @@ -62,6 +63,34 @@ def __init__( self.trajectory = Trajectory(iteration_dir=self.iteration_dir) super().__init__() + def _ranked_dse_candidates(self) -> list[tuple[int, dict[str, str | int | float]]]: + """Return valid trial steps and configurations in descending reward order.""" + tr = self.original_test_run + candidates = [ + row + for row in self.trajectory.dataframe.to_dict(orient="records") + if math.isfinite(row["reward"]) + and all( + math.isfinite(row[f"observation.{metric}"]) + and row[f"observation.{metric}"] != self.rewards.metric_failure + for metric in tr.test.agent_metrics + ) + ] + return [ + (int(row["step"]), {key: row[f"action.{key}"] for key in tr.param_space}) + for row in sorted(candidates, key=lambda row: row["reward"], reverse=True) + ] + + def update_output(self) -> None: + """Publish the DSE recommendation without interrupting execution on output errors.""" + try: + tr = self.original_test_run + output = self.runner.experiment_output + output.update_dse(str(tr.name), tr.param_space, self._ranked_dse_candidates()) + output.write() + except Exception as exc: + logging.warning("Cannot update DSE output for %s: %s", self.original_test_run.name, exc) + @property def upcoming_trial(self) -> int: """ diff --git a/src/cloudai/handlers.py b/src/cloudai/handlers.py new file mode 100644 index 000000000..2563e6e86 --- /dev/null +++ b/src/cloudai/handlers.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-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 argparse +import copy +import logging +import traceback +from pathlib import Path +from typing import Optional + +from cloudai.core import ( + BaseInstaller, + CloudAIGymEnv, + Installable, + Registry, + Runner, + System, + TestScenario, +) +from cloudai.models.scenario import ReportConfig +from cloudai.models.workload import TestDefinition +from cloudai.systems.slurm import SlurmSystem + + +def prepare_installation( + system: System, tests: list[TestDefinition], scenario: Optional[TestScenario] +) -> tuple[list[Installable], BaseInstaller]: + installables: list[Installable] = [] + if scenario: + installables.extend(_scenario_installables(scenario)) + else: + for test in tests: + logging.debug(f"{test.name} has {len(test.installables)} installables.") + installables.extend(test.installables) + + registry = Registry() + installer_class = registry.installers_map.get(system.scheduler, BaseInstaller) + if installer_class is None: + raise NotImplementedError(f"No installer available for scheduler: {system.scheduler}") + installer = installer_class(system) + + return installables, installer + + +def _scenario_installables(scenario: TestScenario) -> list[Installable]: + installables: list[Installable] = [] + for test_run in scenario.test_runs: + logging.debug(f"{test_run.test.name} has {len(test_run.test.installables)} installables.") + installables.extend(test_run.test.installables) + for hook in (test_run.pre_test, test_run.post_test): + if hook is not None: + installables.extend(_scenario_installables(hook)) + return installables + + +def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: + registry = Registry() + + original_test_runs = copy.deepcopy(runner.runner.test_scenario.test_runs) + + has_dependencies = any(tr.dependencies for tr in runner.runner.test_scenario.test_runs) + if has_dependencies: + logging.error( + "Dependencies are not supported for DSE jobs, all cases run consecutively. " + "Please remove dependencies and re-run." + ) + return 1 + + err = 0 + # Capture an unexpected error so reports still generate, then re-raise below. + run_error: Exception | None = None + try: + for tr in runner.runner.test_scenario.test_runs: + test_run = copy.deepcopy(tr) + + agent_type = test_run.test.agent + if not registry.has_agent(agent_type): + logging.error( + f"No agent available for type: {agent_type}. Please make sure {agent_type} " + f"is a valid agent type. Available agents: {registry.agent_names()}" + ) + err = 1 + continue + agent_class = registry.get_agent(agent_type) + + agent_config_data = test_run.test.agent_config or {} + agent_config = agent_class.get_config_class()(**agent_config_data) + env = CloudAIGymEnv( + test_run=test_run, + runner=runner.runner, + rewards=agent_config.rewards, + ) + if agent_config.start_action == "first": + logging.info(f"Using deterministic first sweep for the chosen agent: {env.first_sweep}.") + + agent = agent_class(env, agent_config) + logging.debug(f"Created agent {agent.__class__.__name__}.") + + env.update_output() + try: + err |= agent.run() + finally: + env.update_output() + except Exception as exc: + run_error = exc + logging.exception("DSE job aborted by an unexpected error; generating reports before failing.") + + if args.mode == "run": + runner.runner.test_scenario.test_runs = original_test_runs + generate_reports( + runner.runner.system, + runner.runner.test_scenario, + runner.runner.scenario_root, + error=run_error, + ) + + if run_error is not None: + raise run_error.with_traceback(run_error.__traceback__) + + logging.info("All jobs are complete.") + return err + + +def _record_run_failure(result_dir: Path, error: BaseException) -> None: + """Persist an aborting error into the results dir so the failure is documented with the reports.""" + failure_path = result_dir / "dse_failure.txt" + tb = "".join(traceback.format_exception(type(error), error, error.__traceback__)) + try: + result_dir.mkdir(parents=True, exist_ok=True) + failure_path.write_text(f"DSE job aborted by an unexpected {type(error).__name__}: {error}\n\n{tb}") + logging.info(f"Documented the aborting error in {failure_path}") + except OSError: + logging.exception(f"Failed to write failure report to {failure_path}") + + +def generate_reports( + system: System, + test_scenario: TestScenario, + result_dir: Path, + error: BaseException | None = None, +) -> None: + registry = Registry() + + if error is not None: + _record_run_failure(result_dir, error) + + for name, reporter_class in registry.ordered_scenario_reports(): + logging.debug(f"Generating report '{name}' ({reporter_class.__name__})") + + cfg = registry.report_configs.get(name, ReportConfig(enable=False)) + if scenario_cfg := test_scenario.reports.get(name): + cfg = scenario_cfg + elif isinstance(system, SlurmSystem) and system.reports and name in system.reports: + cfg = system.reports[name] + logging.debug(f"Report '{name}' config is: {cfg.model_dump_json(indent=None)}") + + if not cfg.enable: + logging.debug(f"Skipping report {name} because it is disabled.") + continue + + try: + reporter = reporter_class(system, test_scenario, result_dir, cfg) + reporter.generate() + except Exception as e: + logging.warning(f"Error generating report '{name}', see debug log for details") + logging.debug(e, exc_info=True) + + +def handle_non_dse_job(runner: Runner, args: argparse.Namespace) -> 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 diff --git a/src/cloudai/output.py b/src/cloudai/output.py index c8a39fb5d..ada9aff34 100644 --- a/src/cloudai/output.py +++ b/src/cloudai/output.py @@ -19,9 +19,23 @@ import pathlib import tempfile +import cloudai.metrics import cloudai.models.output +def metric_output(observation: cloudai.metrics.MetricObservation) -> cloudai.models.output.Metric: + """Convert a canonical observation to an output metric.""" + return cloudai.models.output.Metric( + name=observation.metric.display_name, + value=observation.value, + unit=observation.metric.unit, + dimensions=[ + cloudai.models.output.Dimension(name=cloudai.metrics.dimension_label(key), value=str(value)) + for key, value in sorted(observation.dimensions.items()) + ], + ) + + class ExperimentOutput: """Collect experiment results and publish snapshots.""" @@ -56,6 +70,25 @@ def update_test(self, test: cloudai.models.output.Test) -> None: return self.experiment.tests.append(recorded) + def update_dse( + self, + test_id: str, + space: dict[str, list[str | int | float]], + candidates: list[tuple[int, dict[str, str | int | float]]], + ) -> None: + """Store the first ranked candidate with a successful run and its metrics.""" + test = next((test for test in self.experiment.tests if test.id == test_id), None) + if test is None: + raise KeyError(f"Unknown experiment test: {test_id}") + completed_runs = {run.step: run for run in test.runs if run.status == "completed"} + for step, config in candidates: + if step in completed_runs: + test.dse = cloudai.models.output.DSE(space=space, best_step=step, best_config=config) + test.metrics = [metric.model_copy(deep=True) for metric in completed_runs[step].metrics] + return + test.dse = cloudai.models.output.DSE(space=space) + test.metrics = [] + def snapshot(self) -> cloudai.models.output.Experiment: """Return an independent snapshot without finalizing the experiment.""" full = self.experiment.model_copy(deep=True) @@ -87,15 +120,36 @@ def write(self) -> None: logging.warning("Cannot remove temporary experiment output %s: %s", temporary_path, exc) def finish(self, status: cloudai.models.output.Status, finish: datetime.datetime | None) -> None: + for test in self.experiment.tests: + for run in test.runs: + if run.status in ("pending", "running"): + run.status = "unknown" + self._update_test_status(test) + self._update_test_metrics(test) + if status == "completed": + statuses = {test.status for test in self.experiment.tests} + for outcome in ("failed", "cancelled", "unknown"): + if outcome in statuses: + status = outcome + break + for test in self.experiment.tests: + if test.status in ("pending", "running"): + test.status = "completed" if status == "completed" else "unknown" self.experiment.status = status self.experiment.finish = finish - if status == "completed": - for test in self.experiment.tests: - if test.status not in ("failed", "cancelled"): - test.status = "completed" self._update_timing(self.experiment) self.write() + @staticmethod + def _update_test_metrics(test: cloudai.models.output.Test) -> None: + if test.dse is not None: + return + test.metrics = [] + if len(test.runs) == 1: + run = test.runs[0] + if run.status == "completed" and run.step in (None, 0): + test.metrics = [metric.model_copy(deep=True) for metric in run.metrics] + @staticmethod def _update_test_status(test: cloudai.models.output.Test) -> None: statuses = {run.status for run in test.runs} @@ -105,6 +159,10 @@ def _update_test_status(test: cloudai.models.output.Test) -> None: test.status = "cancelled" elif "running" in statuses: test.status = "running" + elif "pending" in statuses: + test.status = "pending" + elif "unknown" in statuses: + test.status = "unknown" elif statuses == {"completed"}: test.status = "completed" diff --git a/src/cloudai/systems/slurm/single_sbatch_runner.py b/src/cloudai/systems/slurm/single_sbatch_runner.py index 5db16fc30..78943933b 100644 --- a/src/cloudai/systems/slurm/single_sbatch_runner.py +++ b/src/cloudai/systems/slurm/single_sbatch_runner.py @@ -21,9 +21,10 @@ from pathlib import Path from typing import Generator, Optional, cast +import cloudai.models.output from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParams -from cloudai.core import BaseJob, Registry, System, TestRun, TestScenario +from cloudai.core import BaseJob, JobStatusResult, Registry, System, TestRun, TestScenario from cloudai.util import format_time_limit, parse_time_limit from .slurm_command_gen_strategy import SlurmCommandGenStrategy @@ -205,10 +206,9 @@ def run(self): finally: self.jobs.remove(job) - self.handle_dse() - self.on_job_completion(job) - self.update_run_output(job) + self.update_run_output(job, JobStatusResult(is_successful=is_completed)) + self.handle_dse() def handle_dse(self): registry = Registry() @@ -221,29 +221,68 @@ def handle_dse(self): agent_config = agent_class.get_config_class()(**agent_config_data) gym = CloudAIGymEnv(tr, self, rewards=agent_config.rewards) - for idx, combination in enumerate(tr.all_combinations, start=1): - sampled_env_params = gym.params.sample(idx) if gym.params is not None else {} - next_tr = tr.apply_params_set(combination, env_params=sampled_env_params) - next_tr.step = idx - next_tr.output_path = self.get_job_output_path(next_tr) - - if not next_tr.test.constraint_check(next_tr, self.system): - continue - - gym.test_run = next_tr - observation = gym.get_observation() - reward = gym.compute_reward(observation) - gym.trajectory.append( - step=idx, - action=combination, - reward=reward, - observation=observation, - env_params=sampled_env_params, - ) + try: + for idx, combination in enumerate(tr.all_combinations, start=1): + sampled_env_params = gym.params.sample(idx) if gym.params is not None else {} + next_tr = tr.apply_params_set(combination, env_params=sampled_env_params) + next_tr.step = idx + next_tr.output_path = self.get_job_output_path(next_tr) + + if not next_tr.test.constraint_check(next_tr, self.system): + continue + + gym.test_run = next_tr + observation = gym.get_observation() + reward = gym.compute_reward(observation) + gym.trajectory.append( + step=idx, + action=combination, + reward=reward, + observation=observation, + env_params=sampled_env_params, + ) + finally: + gym.update_output() def completed_test_runs(self, job: BaseJob) -> list[TestRun]: return list(self.all_trs) + def get_run_output( + self, job: BaseJob, tr: TestRun, result: JobStatusResult | None = None + ) -> cloudai.models.output.Run | None: + run_job = copy.copy(cast(SlurmJob, job)) + allocation_metadata = run_job.metadata + run_job.metadata = None + output_arg = f"--output={tr.output_path.absolute()}/stdout.txt" + steps = ( + [step for step in allocation_metadata.job_steps if output_arg in step.submit_line.split()] + if allocation_metadata is not None + else [] + ) + if result is not None: + if not steps and not (tr.output_path / "stdout.txt").exists(): + result = None + else: + try: + result = tr.test.was_run_successful(tr) + except Exception as exc: + logging.warning("Cannot determine output status for %s: %s", tr.output_path, exc) + return None + run = super().get_run_output(run_job, tr, result) + if run is None: + return None + if run.status == "failed" and any(step.state.startswith("CANCELLED") for step in steps): + run.status = "cancelled" + starts = [self._output_timestamp(step.start_time) for step in steps] + finishes = [self._output_timestamp(step.end_time) for step in steps] + if starts and all(start is not None for start in starts): + run.start = min(start for start in starts if start is not None) + if finishes and all(finish is not None for finish in finishes): + run.finish = max(finish for finish in finishes if finish is not None) + if len(steps) == 1: + run.duration = steps[0].elapsed_time_sec + return run + def _submit_test(self, tr: TestRun) -> SlurmJob: with open(self.scenario_root / "cloudai_sbatch_script.sh", "w") as f: f.write(self.gen_sbatch_content()) diff --git a/src/cloudai/systems/slurm/slurm_job.py b/src/cloudai/systems/slurm/slurm_job.py index c5633b765..2b903e375 100644 --- a/src/cloudai/systems/slurm/slurm_job.py +++ b/src/cloudai/systems/slurm/slurm_job.py @@ -18,9 +18,12 @@ from cloudai.core import BaseJob +from .slurm_metadata import SlurmJobMetadata + @dataclass class SlurmJob(BaseJob): """A job class for execution on a Slurm system.""" nodes: list[str] = field(default_factory=list, init=False) + metadata: SlurmJobMetadata | None = field(default=None, init=False) diff --git a/src/cloudai/systems/slurm/slurm_runner.py b/src/cloudai/systems/slurm/slurm_runner.py index 93a177962..4928cabb2 100644 --- a/src/cloudai/systems/slurm/slurm_runner.py +++ b/src/cloudai/systems/slurm/slurm_runner.py @@ -14,13 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime import logging from pathlib import Path from typing import cast import toml -from cloudai.core import BaseJob, BaseRunner, System, TestRun, TestScenario +import cloudai.models.output +import cloudai.output +from cloudai.core import BaseJob, BaseRunner, JobStatusResult, System, TestRun, TestScenario from .slurm_command_gen_strategy import SlurmCommandGenStrategy from .slurm_job import SlurmJob @@ -36,6 +39,44 @@ def __init__(self, mode: str, system: System, test_scenario: TestScenario, outpu self.system = cast(SlurmSystem, system) self.pinned_nodes: dict[str, list[str]] = {} + def get_run_output( + self, job: BaseJob, tr: TestRun, result: JobStatusResult | None = None + ) -> cloudai.models.output.Run | None: + metadata = cast(SlurmJob, job).metadata + status: cloudai.models.output.Status = "pending" + metrics: list[cloudai.models.output.Metric] = [] + if result is not None: + status = "completed" if result.is_successful else "failed" + if job.terminated_by_dependency or (metadata is not None and metadata.state.startswith("CANCELLED")): + status = "cancelled" + if status == "completed": + try: + metrics = [ + cloudai.output.metric_output(observation) + for observation in tr.test.metric_observations(self.system, tr) + ] + except Exception as exc: + logging.warning("Cannot extract output metrics for Slurm job %s: %s", job.id, exc) + return cloudai.models.output.Run( + path=str(tr.output_path.absolute()), + jobid=str(job.id), + status=status, + metrics=metrics, + start=self._output_timestamp(metadata.start_time) if metadata is not None else None, + finish=self._output_timestamp(metadata.end_time) if metadata is not None else None, + duration=metadata.elapsed_time_sec if metadata is not None else None, + iteration=tr.current_iteration, + step=tr.step, + ) + + @staticmethod + def _output_timestamp(value: str) -> datetime.datetime | None: + try: + timestamp = datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return timestamp if timestamp.utcoffset() is not None else None + def submit_test(self, tr: TestRun) -> None: if tr.pin_nodes and tr.name in self.pinned_nodes: tr.nodes = self.pinned_nodes[tr.name].copy() @@ -119,7 +160,11 @@ def _get_job_metadata( def store_job_metadata(self, job: SlurmJob): system = cast(SlurmSystem, self.system) steps_metadata = [self._mock_job_metadata()] if self.mode == "dry-run" else system.get_job_status(job) + if not steps_metadata: + logging.warning("No Slurm accounting metadata available for job %s", job.id) + return slurm_job_file, job_meta = self._get_job_metadata(job, steps_metadata) + job.metadata = job_meta logging.debug(f"Storing job metadata for job {job.id} to {slurm_job_file}") with slurm_job_file.open("w") as job_file: diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 7b3d31010..8b7fb4aec 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -481,6 +481,7 @@ def get_job_status(self, job: BaseJob, retry_threshold: int = 3) -> list[SlurmSt retry_count = 0 command = ( + "TZ=UTC SLURM_TIME_FORMAT='%Y-%m-%dT%H:%M:%SZ' " f"sacct -j {job.id} --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine " "--delimiter='|' -p --noheader" ) diff --git a/src/cloudai/systems/standalone/standalone_runner.py b/src/cloudai/systems/standalone/standalone_runner.py index 4e03ef0a8..fe2d6b9db 100644 --- a/src/cloudai/systems/standalone/standalone_runner.py +++ b/src/cloudai/systems/standalone/standalone_runner.py @@ -19,8 +19,8 @@ from pathlib import Path from typing import cast -import cloudai.metrics import cloudai.models.output +import cloudai.output from cloudai.core import BaseJob, BaseRunner, JobIdRetrievalError, JobStatusResult, System, TestRun, TestScenario from cloudai.util import CommandShell @@ -52,7 +52,7 @@ def get_run_output( if result.is_successful: try: observations = tr.test.metric_observations(self.system, tr) - metrics = [self._metric_output(observation) for observation in observations] + metrics = [cloudai.output.metric_output(observation) for observation in observations] except Exception as exc: logging.warning("Cannot extract output metrics for standalone job %s: %s", job.id, exc) return cloudai.models.output.Run( @@ -70,22 +70,6 @@ def on_job_completion(self, job: BaseJob) -> None: standalone_job = cast(StandaloneJob, job) standalone_job.finish = datetime.datetime.now(datetime.timezone.utc) - @staticmethod - def _metric_output(observation: cloudai.metrics.MetricObservation) -> cloudai.models.output.Metric: - dimensions = [ - cloudai.models.output.Dimension( - name=cloudai.metrics.dimension_label(key), - value=str(value), - ) - for key, value in sorted(observation.dimensions.items()) - ] - return cloudai.models.output.Metric( - name=observation.metric.display_name, - value=observation.value, - unit=observation.metric.unit, - dimensions=dimensions, - ) - def _submit_test(self, tr: TestRun) -> StandaloneJob: logging.info(f"Running test: {tr.name}") tr.output_path = self.get_job_output_path(tr) diff --git a/tests/systems/slurm/test_runner.py b/tests/systems/slurm/test_runner.py new file mode 100644 index 000000000..e065d5a3b --- /dev/null +++ b/tests/systems/slurm/test_runner.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import datetime +import pathlib +from unittest import mock + +import pytest + +import cloudai.core +import cloudai.metrics +from cloudai.systems.slurm import SlurmJob, SlurmRunner, SlurmSystem +from cloudai.systems.slurm.slurm_metadata import SlurmStepMetadata + + +@pytest.mark.parametrize( + "successful,state,status", + [(True, "FAILED", "completed"), (False, "COMPLETED", "failed"), (False, "CANCELLED", "cancelled")], +) +@pytest.mark.parametrize("step", [0, 2]) +def test_slurm_run_output( + tmp_path: pathlib.Path, + base_tr: cloudai.core.TestRun, + slurm_system: SlurmSystem, + successful: bool, + state: str, + status: str, + step: int, +) -> None: + runner = SlurmRunner("run", slurm_system, cloudai.core.TestScenario(name="scenario", test_runs=[base_tr]), tmp_path) + base_tr.output_path.mkdir(parents=True) + base_tr.step = step + job = SlurmJob(base_tr, id=123) + runner.update_run_output(job) + metadata = SlurmStepMetadata( + job_id=123, + step_id="", + name="job", + state=state, + exit_code="1:0", + elapsed_time_sec=3, + start_time="2026-01-02T03:04:05Z", + end_time="2026-01-02T03:04:08Z", + submit_line="sbatch run.sh", + ) + observation = cloudai.metrics.MetricObservation(cloudai.metrics.BANDWIDTH, 12.5, {"size_bytes": 1024}) + with ( + mock.patch.object(SlurmSystem, "get_job_status", return_value=[metadata]) as get_metadata, + mock.patch.object( + runner, + "get_cmd_gen_strategy", + return_value=mock.Mock(gen_srun_command=lambda: "srun cmd", generate_test_command=lambda: ["cmd"]), + ), + mock.patch.object( + cloudai.core.TestDefinition, + "was_run_successful", + return_value=cloudai.core.JobStatusResult(is_successful=successful), + ), + mock.patch.object( + cloudai.core.TestDefinition, "metric_observations", return_value=[observation] + ) as get_metrics, + ): + runner.store_job_metadata(job) + runner.update_run_output(job, runner.get_job_status(job)) + get_metadata.assert_called_once_with(job) + assert get_metrics.call_count == int(successful) + + base_tr.step = 3 + runner.shutting_down = status == "cancelled" + runner.finish_output(successful=True) + experiment = runner.experiment_output.snapshot() + assert experiment.status == status + test = experiment.tests[0] + assert test.metrics == (test.runs[0].metrics if successful and step == 0 else []) + start = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc) + assert [run.model_dump() for run in test.runs] == [ + { + "path": str(base_tr.output_path.absolute()), + "jobid": "123", + "status": status, + "metrics": [ + { + "name": "Bandwidth", + "value": 12.5, + "unit": "GB/s", + "dimensions": [{"name": "Size", "value": "1024", "unit": "", "is_x": False}], + } + ] + if successful + else [], + "start": start, + "finish": start + datetime.timedelta(seconds=3), + "duration": 3, + "iteration": 0, + "step": step, + } + ] + + +@pytest.mark.parametrize("timestamp", ["Unknown", "", "2026-01-02T03:04:05"]) +def test_slurm_output_unknown_timing_and_metric_failure( + tmp_path: pathlib.Path, + base_tr: cloudai.core.TestRun, + slurm_system: SlurmSystem, + timestamp: str, + caplog: pytest.LogCaptureFixture, +) -> None: + runner = SlurmRunner("run", slurm_system, cloudai.core.TestScenario(name="scenario", test_runs=[base_tr]), tmp_path) + job = SlurmJob(base_tr, id=123) + with mock.patch.object(SlurmSystem, "get_job_status", return_value=[]): + runner.store_job_metadata(job) + with mock.patch.object( + cloudai.core.TestDefinition, "metric_observations", side_effect=ValueError("broken metrics") + ): + run = runner.get_run_output(job, base_tr, cloudai.core.JobStatusResult(is_successful=True)) + assert run is not None + assert run.model_dump() == { + "path": str(base_tr.output_path.absolute()), + "jobid": "123", + "status": "completed", + "metrics": [], + "start": None, + "finish": None, + "duration": None, + "iteration": 0, + "step": 0, + } + assert runner._output_timestamp(timestamp) is None + assert "broken metrics" in caplog.text diff --git a/tests/systems/slurm/test_system.py b/tests/systems/slurm/test_system.py index 561f742ef..2d694fd87 100644 --- a/tests/systems/slurm/test_system.py +++ b/tests/systems/slurm/test_system.py @@ -902,6 +902,11 @@ def test_get_job_status(slurm_system: SlurmSystem, stdout: str, stderr: str, exp slurm_system.get_job_status(job) else: assert slurm_system.get_job_status(job) == expected + slurm_system.cmd_shell.execute.assert_called_with( + "TZ=UTC SLURM_TIME_FORMAT='%Y-%m-%dT%H:%M:%SZ' " + "sacct -j 1 --format=JobID,JobName,State,ExitCode,Start,End,ElapsedRAW,SubmitLine " + "--delimiter='|' -p --noheader" + ) sacct_output = """2623913,job,COMPLETED,0:0,2025-05-09T01:34:52,2025-05-09T01:59:27,1475,sbatch sbatch_script.sh, diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 53105c5d4..bf6202bb5 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -27,7 +27,7 @@ import cloudai.models.output from cloudai.cli import setup_logging -from cloudai.cli.handlers import handle_dry_run_and_run +from cloudai.cli.cli import handle_dry_run_and_run from cloudai.core import CommandGenStrategy, GitRepo, TestDefinition, TestRun, TestScenario from cloudai.models.scenario import TestRunDetails from cloudai.systems.slurm import SlurmCommandGenStrategy, SlurmRunner, SlurmSystem diff --git a/tests/test_handlers.py b/tests/test_handlers.py index 560b2f409..9f401ec73 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -17,22 +17,17 @@ import argparse import copy from pathlib import Path -from typing import Any, ClassVar, Iterator, Optional +from typing import Any, ClassVar, Iterator, Optional, cast from unittest.mock import MagicMock import pandas as pd import pytest from pydantic import Field -from cloudai.cli.handlers import ( - handle_dse_job, - prepare_installation, - validate_domain_randomization_active, - verify_system_configs, - verify_test_configs, - verify_test_scenarios, -) -from cloudai.configurator.env_params import EnvParamSpec +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.core import ( BaseAgent, BaseAgentConfig, @@ -46,6 +41,10 @@ TestScenario, TestScenarioParsingError, ) +from cloudai.handlers import ( + handle_dse_job, + prepare_installation, +) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import CmdArgs, TestDefinition from cloudai.reporter import StatusReporter @@ -397,6 +396,86 @@ def test_handle_dse_job_invokes_agent_run( assert CustomRunStubAgent.run_calls == 1 +@pytest.mark.parametrize( + "failed_steps,last_reward,last_observation,best_step", + [ + ([], 2.0, 2.0, 2), + ([2], 2.0, 2.0, 1), + ([1, 2], 2.0, 2.0, None), + ([], float("inf"), 2.0, 1), + ([], 2.0, float("nan"), 1), + ([], 2.0, -1.0, 1), + ([], 1.0, 2.0, 1), + ], +) +def test_dse_output_selects_completed_trial( + slurm_system: SlurmSystem, + dse_tr: TestRun, + custom_run_agent_name: str, + monkeypatch: pytest.MonkeyPatch, + failed_steps: list[int], + last_reward: float, + last_observation: float, + best_step: int | None, +) -> None: + dse_tr.test.agent = custom_run_agent_name + runner = Runner("dry-run", slurm_system, TestScenario(name="scenario", test_runs=[dse_tr])) + + def run(agent: CustomRunStubAgent) -> int: + env = cast(CloudAIGymEnv, agent.env) + test = env.runner.experiment_output.snapshot().tests[0] + assert test.dse is not None + assert test.dse.model_dump() == { + "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, + "best_step": None, + "best_config": None, + } + for step, value in enumerate(["value1", "value2"], start=1): + env.runner.experiment_output.update_run( + dse_tr.name, + cloudai.models.output.Run( + path=str(env.iteration_dir / str(step)), + jobid=str(step), + step=step, + status="failed" if step in failed_steps else "completed", + metrics=[cloudai.models.output.Metric(name="Bandwidth", value=12.5 * step, unit="GB/s")], + ), + ) + env.trajectory.append( + step=step, + action={"extra_env_vars.VAR1": value}, + reward=last_reward if step == 2 else 1.0, + observation={metric: last_observation if step == 2 else 1.0 for metric in dse_tr.test.agent_metrics}, + env_params={}, + ) + if failed_steps: + raise RuntimeError("trial failed") + return 0 + + monkeypatch.setattr(CustomRunStubAgent, "run", run) + if failed_steps: + with pytest.raises(RuntimeError, match="trial failed"): + handle_dse_job(runner, argparse.Namespace(mode="dry-run")) + else: + assert handle_dse_job(runner, argparse.Namespace(mode="dry-run")) == 0 + + stored = cloudai.models.output.Experiment.model_validate_json( + (runner.runner.scenario_root / "experiment.json").read_text() + ) + assert stored.tests[0].dse is not None + assert stored.tests[0].dse.model_dump() == { + "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, + "best_step": best_step, + "best_config": {"extra_env_vars.VAR1": f"value{best_step}"} if best_step is not None else None, + } + assert [metric.model_dump() for metric in stored.tests[0].metrics] == ( + [{"name": "Bandwidth", "value": 12.5 * best_step, "unit": "GB/s", "dimensions": []}] + if best_step is not None + else [] + ) + assert len(stored.tests[0].runs) == 2 + + def test_handle_dse_job_propagates_agent_run_nonzero_rc( slurm_system: SlurmSystem, dse_tr: TestRun, diff --git a/tests/test_output.py b/tests/test_output.py index 93622f5b7..65e74916f 100644 --- a/tests/test_output.py +++ b/tests/test_output.py @@ -17,11 +17,16 @@ import datetime import pathlib +import pytest + import cloudai.models.output import cloudai.output -def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathlib.Path) -> None: +@pytest.mark.parametrize("status", ["completed", "failed", "cancelled"]) +def test_experiment_output_preserves_runs_and_finalizes_failure( + tmp_path: pathlib.Path, status: cloudai.models.output.Status +) -> None: start = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc) experiment = cloudai.models.output.Experiment( id="experiment", @@ -29,7 +34,10 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli status="running", path=str(tmp_path), start=start, - tests=[cloudai.models.output.Test(id="case", name="workload", path=str(tmp_path / "case"))], + tests=[ + cloudai.models.output.Test(id=case, name="workload", path=str(tmp_path / case)) + for case in ("case", "interrupted", "not-started") + ], ) experiment_output = cloudai.output.ExperimentOutput(experiment, tmp_path) first_run = cloudai.models.output.Run( @@ -44,7 +52,7 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli second_run = cloudai.models.output.Run( path=str(tmp_path / "case" / "1"), jobid="102", - status="running", + status="pending", start=start + datetime.timedelta(seconds=2), iteration=1, step=0, @@ -52,17 +60,22 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli experiment_output.update_run("case", first_run) experiment_output.update_run("case", second_run) + assert experiment_output.snapshot().tests[0].status == "pending" second_run.status = "failed" second_run.finish = start + datetime.timedelta(seconds=4) experiment_output.update_run("case", second_run) - experiment_output.finish("failed", start + datetime.timedelta(seconds=5)) + experiment_output.update_run( + "interrupted", + cloudai.models.output.Run(path=str(tmp_path / "interrupted" / "0"), jobid="103", status="pending"), + ) + experiment_output.finish(status, start + datetime.timedelta(seconds=5)) stored = cloudai.models.output.Experiment.model_validate_json((tmp_path / "experiment.json").read_text()) assert stored.model_dump() == { "id": "experiment", "name": "scenario", "description": None, - "status": "failed", + "status": "cancelled" if status == "cancelled" else "failed", "path": str(tmp_path), "start": start, "finish": start + datetime.timedelta(seconds=5), @@ -100,6 +113,38 @@ def test_experiment_output_preserves_runs_and_finalizes_failure(tmp_path: pathli }, ], "dse": None, - } + }, + { + "id": "interrupted", + "name": "workload", + "description": None, + "status": "unknown", + "path": str(tmp_path / "interrupted"), + "metrics": [], + "runs": [ + { + "path": str(tmp_path / "interrupted" / "0"), + "jobid": "103", + "status": "unknown", + "metrics": [], + "start": None, + "finish": None, + "duration": None, + "iteration": None, + "step": None, + } + ], + "dse": None, + }, + { + "id": "not-started", + "name": "workload", + "description": None, + "status": "unknown", + "path": str(tmp_path / "not-started"), + "metrics": [], + "runs": [], + "dse": None, + }, ], } diff --git a/tests/test_reporter.py b/tests/test_reporter.py index 045eb50d6..b6e3a0f65 100644 --- a/tests/test_reporter.py +++ b/tests/test_reporter.py @@ -26,8 +26,8 @@ import toml from cloudai import TestRun, TestScenario -from cloudai.cli.handlers import generate_reports from cloudai.core import CommandGenStrategy, Registry, Reporter, System +from cloudai.handlers import generate_reports from cloudai.models.scenario import ReportConfig, TestRunDetails from cloudai.report_generator.dse_report import build_dse_summaries from cloudai.reporter import DSEReporter, JUnitReporter, PerTestReporter, ReportItem, StatusReporter, TarballReporter diff --git a/tests/test_single_sbatch_runner.py b/tests/test_single_sbatch_runner.py index 4afe91b96..5464453c3 100644 --- a/tests/test_single_sbatch_runner.py +++ b/tests/test_single_sbatch_runner.py @@ -15,6 +15,7 @@ # limitations under the License. import copy +import datetime import re from pathlib import Path from typing import Generator, Optional, cast @@ -24,9 +25,13 @@ import pytest import toml +import cloudai.metrics +import cloudai.models.output +from cloudai.configurator import CloudAIGymEnv from cloudai.configurator.env_params import EnvParams, EnvParamSpec -from cloudai.core import Registry, System, TestRun, TestScenario +from cloudai.core import JobStatusResult, Registry, System, TestRun, TestScenario from cloudai.systems.slurm import SingleSbatchRunner, SlurmJob, SlurmJobMetadata, SlurmSystem +from cloudai.systems.slurm.slurm_metadata import SlurmStepMetadata from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition from cloudai.workloads.nccl_test.slurm_command_gen_strategy import NcclTestSlurmCommandGenStrategy from cloudai.workloads.sleep import SleepCmdArgs, SleepTestDefinition @@ -582,22 +587,108 @@ def submit_after_shutdown(_: TestRun) -> SlurmJob: assert runner.jobs == [] -def test_run_removes_completed_job_from_tracking(sleep_tr: TestRun, slurm_system: SlurmSystem) -> None: - tc = TestScenario(name="tc", test_runs=[sleep_tr]) +@pytest.mark.parametrize("status", ["completed", "failed", "cancelled", "unknown"]) +def test_run_removes_completed_job_from_tracking(sleep_tr: TestRun, slurm_system: SlurmSystem, status: str) -> None: + second_tr = copy.deepcopy(sleep_tr) + second_tr.name = "second" + tc = TestScenario(name="tc", test_runs=[sleep_tr, second_tr]) runner = SingleSbatchRunner(mode="run", system=slurm_system, test_scenario=tc, output_path=slurm_system.output_path) + runs = list(runner.all_trs) job = SlurmJob(sleep_tr, id=123) runner._submit_test = Mock(return_value=job) runner.handle_dse = Mock() - runner.on_job_completion = Mock() - + runner.on_job_completion = Mock(side_effect=runner.store_job_metadata) + allocation = SlurmStepMetadata( + job_id=123, + step_id="", + name="allocation", + state="CANCELLED", + exit_code="1:0", + start_time="2026-01-02T03:04:00Z", + end_time="2026-01-02T03:04:30Z", + elapsed_time_sec=30, + submit_line="sbatch run.sh", + ) + steps = [allocation] + for tr, step, start, end in [(sleep_tr, "3", 1, 4), (second_tr, "4", 5, 9), (second_tr, "5", 7, 12)]: + if tr is second_tr and status == "unknown": + continue + steps.append( + allocation.model_copy( + update={ + "step_id": step, + "submit_line": f"srun --output={tr.output_path.absolute()}/stdout.txt bash run.sh", + "state": "CANCELLED" if tr is second_tr and status == "cancelled" else "FAILED", + "start_time": f"2026-01-02T03:04:{start:02}Z", + "end_time": f"2026-01-02T03:04:{end:02}Z", + "elapsed_time_sec": end - start, + } + ) + ) with ( patch.object(SlurmSystem, "is_job_completed", return_value=True), + patch.object(SlurmSystem, "get_job_status", return_value=steps), patch.object(SlurmSystem, "kill") as kill, + patch.object( + SleepTestDefinition, + "was_run_successful", + side_effect=lambda tr: JobStatusResult( + is_successful=tr.name == sleep_tr.name or status == "completed", + ), + ), + patch.object( + SleepTestDefinition, + "metric_observations", + side_effect=lambda system, tr: [ + cloudai.metrics.MetricObservation(cloudai.metrics.BANDWIDTH, 5 if tr.name == sleep_tr.name else 10, {}), + ], + ), ): runner.run() + runner.finish_output(successful=True) kill.assert_not_called() assert runner.jobs == [] + stored = cloudai.models.output.Experiment.model_validate_json( + (runner.scenario_root / "experiment.json").read_text() + ) + assert stored.status == status + start = datetime.datetime(2026, 1, 2, 3, 4, tzinfo=datetime.timezone.utc) + expected = [] + for index, tr in enumerate(runs): + run_status = "completed" if index == 0 else status + metrics = [{"name": "Bandwidth", "value": 5 * (index + 1), "unit": "GB/s", "dimensions": []}] + if run_status != "completed": + metrics = [] + expected.append( + { + "id": tr.name, + "name": "sleep", + "description": "desc", + "status": run_status, + "path": str(runner.scenario_root / tr.name), + "metrics": metrics, + "dse": None, + "runs": [ + { + "path": str(tr.output_path.absolute()), + "jobid": "123", + "status": run_status, + "metrics": metrics, + "start": start + datetime.timedelta(seconds=1 if index == 0 else 5) + if run_status != "unknown" + else None, + "finish": start + datetime.timedelta(seconds=4 if index == 0 else 12) + if run_status != "unknown" + else None, + "duration": (3 if index == 0 else 7) if run_status != "unknown" else None, + "iteration": 0, + "step": 0, + } + ], + } + ) + assert [test.model_dump() for test in stored.tests] == expected def test_pre_test(nccl_tr: TestRun, sleep_tr: TestRun, slurm_system: SlurmSystem) -> None: @@ -624,7 +715,8 @@ def test_pre_test(nccl_tr: TestRun, sleep_tr: TestRun, slurm_system: SlurmSystem ) -def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None: +@pytest.mark.parametrize("failed_step", [None, 2]) +def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem, failed_step: int | None) -> None: tc = TestScenario(name="tc", test_runs=[dse_tr]) runner = SingleSbatchRunner(mode="run", system=slurm_system, test_scenario=tc, output_path=slurm_system.output_path) dse_tr.output_path = slurm_system.output_path / dse_tr.name @@ -632,7 +724,40 @@ def test_trajectory_saved(dse_tr: TestRun, slurm_system: SlurmSystem) -> None: trajectory_path = runner.scenario_root / dse_tr.name / f"{dse_tr.current_iteration}" / "trajectory.csv" trajectory_path.unlink(missing_ok=True) - runner.handle_dse() + dse_tr.test.agent_reward_function = "identity" + for tr in runner.all_trs: + runner.experiment_output.update_run( + dse_tr.name, + cloudai.models.output.Run( + path=str(tr.output_path), + jobid="123", + iteration=0, + step=tr.step, + status="failed" if tr.step == failed_step else "completed", + metrics=[cloudai.models.output.Metric(name="Bandwidth", value=tr.step * 12.5, unit="GB/s")], + ), + ) + with patch.object( + CloudAIGymEnv, + "get_observation", + side_effect=[{metric: float(step) for metric in dse_tr.test.agent_metrics} for step in (1, 2)], + ): + runner.handle_dse() + stored = cloudai.models.output.Experiment.model_validate_json( + (runner.scenario_root / "experiment.json").read_text() + ) + test = stored.tests[0] + best_step = 1 if failed_step else 2 + assert test.dse is not None + assert test.dse.model_dump() == { + "space": {"extra_env_vars.VAR1": ["value1", "value2"]}, + "best_step": best_step, + "best_config": {"extra_env_vars.VAR1": f"value{best_step}"}, + } + assert [metric.model_dump() for metric in test.metrics] == [ + {"name": "Bandwidth", "value": best_step * 12.5, "unit": "GB/s", "dimensions": []}, + ] + assert len(test.runs) == 2 assert trajectory_path.exists() df = pd.read_csv(trajectory_path)