diff --git a/src/cloudai/cli/cli.py b/src/cloudai/cli/cli.py index 6df993218..2fe161144 100644 --- a/src/cloudai/cli/cli.py +++ b/src/cloudai/cli/cli.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -104,6 +104,13 @@ def common_options(f): type=click.Path(exists=True, resolve_path=True, path_type=Path), help="Scenario config path.", )(f) + f = click.option( + "--hook-dir", + required=False, + default=None, + type=click.Path(exists=True, resolve_path=True, path_type=Path, file_okay=False, dir_okay=True), + help="Directory with hook scenario and test configs.", + )(f) return f @@ -136,18 +143,28 @@ def main(log_file, log_level): @main.command() @common_options -def install(system_cfg: Path, tests_dir: Path, scenario_cfg: Path): +def install(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_dir: Path | None): """Install the necessary components for workloads.""" - args = argparse.Namespace(system_config=system_cfg, tests_dir=tests_dir, test_scenario=scenario_cfg, mode="install") + args = argparse.Namespace( + system_config=system_cfg, + tests_dir=tests_dir, + test_scenario=scenario_cfg, + hook_dir=hook_dir, + mode="install", + ) exit(handle_install_and_uninstall(args)) @main.command() @common_options -def uninstall(system_cfg: Path, tests_dir: Path, scenario_cfg: Path): +def uninstall(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_dir: Path | None): """Uninstall the components used by workloads.""" args = argparse.Namespace( - system_config=system_cfg, tests_dir=tests_dir, test_scenario=scenario_cfg, mode="uninstall" + system_config=system_cfg, + tests_dir=tests_dir, + test_scenario=scenario_cfg, + hook_dir=hook_dir, + mode="uninstall", ) exit(handle_install_and_uninstall(args)) @@ -161,6 +178,7 @@ def dry_run( system_cfg: Path, tests_dir: Path, scenario_cfg: Path, + hook_dir: Path | None, output_dir: Path, enable_cache_without_check: bool, single_sbatch: bool, @@ -170,6 +188,7 @@ def dry_run( system_config=system_cfg, tests_dir=tests_dir, test_scenario=scenario_cfg, + hook_dir=hook_dir, output_dir=output_dir, mode="dry-run", enable_cache_without_check=enable_cache_without_check, @@ -187,6 +206,7 @@ def run( system_cfg: Path, tests_dir: Path, scenario_cfg: Path, + hook_dir: Path | None, output_dir: Path, enable_cache_without_check: bool, single_sbatch: bool, @@ -200,6 +220,7 @@ def run( system_config=system_cfg, tests_dir=tests_dir, test_scenario=scenario_cfg, + hook_dir=hook_dir, output_dir=output_dir, mode="run", enable_cache_without_check=enable_cache_without_check, @@ -216,7 +237,7 @@ def run( type=click.Path(exists=True, resolve_path=True, path_type=Path, file_okay=False), help="Path to a scenario results directory.", ) -def generate_report(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, result_dir: Path): +def generate_report(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, hook_dir: Path | None, result_dir: Path): """ Generate a report from the results of a scenario. @@ -224,7 +245,11 @@ def generate_report(system_cfg: Path, tests_dir: Path, scenario_cfg: Path, resul command. """ args = argparse.Namespace( - system_config=system_cfg, tests_dir=tests_dir, test_scenario=scenario_cfg, result_dir=result_dir + system_config=system_cfg, + tests_dir=tests_dir, + test_scenario=scenario_cfg, + hook_dir=hook_dir, + result_dir=result_dir, ) exit(handle_generate_report(args)) diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index 3f6488f86..09de12f85 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -64,7 +64,7 @@ def handle_install_and_uninstall(args: argparse.Namespace) -> int: Args: args (argparse.Namespace): The parsed command-line arguments. """ - parser = Parser(args.system_config) + parser = Parser(args.system_config, args.hook_dir or HOOK_ROOT) system, tests, scenario = parser.parse(args.tests_dir, args.test_scenario) system.update() @@ -251,7 +251,7 @@ def register_signal_handlers(signal_handler: Callable) -> None: def _setup_system_and_scenario( args: argparse.Namespace, ) -> tuple[System, TestScenario, list[TestDefinition]] | None: - parser = Parser(args.system_config) + 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: @@ -362,7 +362,7 @@ def handle_generate_report(args: argparse.Namespace) -> int: Args: args (argparse.Namespace): The parsed command-line arguments. """ - parser = Parser(args.system_config) + 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 diff --git a/src/cloudai/parser.py b/src/cloudai/parser.py index 5dbe22c0b..890709d8a 100644 --- a/src/cloudai/parser.py +++ b/src/cloudai/parser.py @@ -37,21 +37,23 @@ from .toml_utils import format_toml_decode_error HOOK_ROOT = Path("conf/hook") -HOOK_TEST_ROOT = HOOK_ROOT / "test" class Parser: """Main parser for parsing all types of configurations.""" - def __init__(self, system_config_path: Path) -> None: + def __init__(self, system_config_path: Path, hook_root: Path = HOOK_ROOT) -> 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. """ logging.debug(f"Initializing parser with: {system_config_path=}") self.system_config_path = system_config_path + self.hook_root = hook_root + self.hook_test_root = hook_root / "test" self._system: Optional[System] = None @property @@ -90,12 +92,14 @@ def parse( except TestConfigParsingError: exit(1) # exit right away to keep error message readable for users - if not HOOK_ROOT.exists(): - logging.debug(f"HOOK_ROOT path '{HOOK_ROOT}' does not exist.") + 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(HOOK_TEST_ROOT.glob("*.toml")), self.system) if HOOK_TEST_ROOT.exists() else [] + 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 @@ -106,10 +110,10 @@ def parse( test_mapping = {t.name: t for t in tests} hook_test_scenario_mapping = {} - if HOOK_ROOT.exists() and list(HOOK_ROOT.glob("*.toml")): + if self.hook_root.exists() and list(self.hook_root.glob("*.toml")): try: hook_test_scenario_mapping = self.parse_hooks( - list(HOOK_ROOT.glob("*.toml")), self.system, {t.name: t for t in hook_tests} + 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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 2a3f394ec..7eda4dc2d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -46,6 +46,16 @@ def test_tests_dir_is_optional(tmp_path: Path): assert "Missing option '--tests-dir'" not in result.output +@pytest.mark.parametrize("subcommand", ["install", "uninstall", "dry-run", "run", "generate-report"]) +def test_hook_dir_is_available_on_run_commands(subcommand: str): + runner = CliRunner() + + result = runner.invoke(main, [subcommand, "--help"]) + + assert result.exit_code == 0 + assert "--hook-dir DIRECTORY" in result.output + + @pytest.mark.parametrize( "subcommand", ["dry-run", "generate-report", "install", "list", "run", "uninstall", "verify-configs"] ) diff --git a/tests/test_parser.py b/tests/test_parser.py index 63a9413da..17e1c11c9 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -39,6 +39,75 @@ def test_no_tests_dir(self, parser: Parser): parser.parse(tests_dir, None) assert "Test path" in str(exc_info.value) + def test_custom_hook_root_is_used(self, parser: Parser, tmp_path: Path): + hook_root = tmp_path / "hooks" + parser = Parser(parser.system_config_path, hook_root) + + assert parser.hook_root == hook_root + assert parser.hook_test_root == hook_root / "test" + + def test_custom_hook_root_is_used_by_parse(self, parser: Parser, tmp_path: Path): + hook_root = tmp_path / "hooks" + hook_test_root = hook_root / "test" + hook_test_root.mkdir(parents=True) + + (hook_test_root / "custom_hook_test.toml").write_text( + 'name = "custom_hook_test"\n' + 'description = "custom hook test"\n' + 'test_template_name = "Sleep"\n' + "\n" + "[cmd_args]\n" + "seconds = 1\n" + ) + + parser = Parser(parser.system_config_path, hook_root) + + _, tests, _ = parser.parse(None, None) + + assert "custom_hook_test" in {test.name for test in tests} + + def test_custom_hook_root_is_used_for_hook_scenario_resolution(self, parser: Parser, tmp_path: Path): + """A hook *scenario* toml (referenced via pre_test) must also be resolved from the custom hook_root, + not just hook test tomls under `/test`.""" + hook_root = tmp_path / "hooks" + hook_test_root = hook_root / "test" + hook_test_root.mkdir(parents=True) + + (hook_test_root / "custom_hook_test.toml").write_text( + 'name = "custom_hook_test"\n' + 'description = "custom hook test"\n' + 'test_template_name = "Sleep"\n' + "\n" + "[cmd_args]\n" + "seconds = 1\n" + ) + (hook_root / "custom_hook_scenario.toml").write_text( + 'name = "custom_hook"\n\n[[Tests]]\nid = "Tests.hook"\ntest_name = "custom_hook_test"\n' + ) + + tests_dir = tmp_path / "tests" + tests_dir.mkdir() + (tests_dir / "main_test.toml").write_text( + 'name = "main_test"\ndescription = "main test"\ntest_template_name = "Sleep"\n\n[cmd_args]\nseconds = 1\n' + ) + test_scenario_path = tmp_path / "test_scenario.toml" + test_scenario_path.write_text( + 'name = "main-scenario"\n' + 'pre_test = "custom_hook"\n\n' + "[[Tests]]\n" + 'id = "Tests.main"\n' + 'test_name = "main_test"\n' + ) + + parser = Parser(parser.system_config_path, hook_root) + + _, _, test_scenario = parser.parse(tests_dir, test_scenario_path) + + assert test_scenario is not None + pre_test = test_scenario.test_runs[0].pre_test + assert pre_test is not None + assert pre_test.name == "custom_hook" + @patch("cloudai.test_parser.TestParser.parse_all") def test_no_scenario(self, test_parser: Mock, parser: Parser): tests_dir = parser.system_config_path.parent.parent / "test"