From ffcf00719af8a5a3589dfe0354ca33ef0f0595b2 Mon Sep 17 00:00:00 2001 From: shreyaskommuri Date: Mon, 29 Jun 2026 10:41:30 -0700 Subject: [PATCH 1/3] Allow external callers to locate hook configs Expose an optional --hook-dir across scenario commands and pass the resolved directory into Parser so tools such as CloudAI Autotune can invoke hook-bearing scenarios outside the CloudAI repository. Constraint: Preserve conf/hook as the parser default and keep hook-free external invocations backward compatible. Rejected: Require callers to change into the CloudAI checkout | output paths and orchestration should not depend on process working directory. Confidence: high Scope-risk: narrow Directive: Keep scenario config roots explicit at the CLI boundary when external orchestration depends on them. Tested: uv run pytest -q (1636 passed, 4 skipped); uv run pytest -q -m ci_only (490 passed); uv run pytest --dead-fixtures; uv run pre-commit run --all-files; external NCCL dry-run from CloudAI-Autotune cwd exits 0 with --hook-dir Not-tested: Live NCCL execution on a Slurm cluster. Signed-off-by: shreyaskommuri --- src/cloudai/cli/cli.py | 39 ++++++++++++++++++++++++++++++------- src/cloudai/cli/handlers.py | 6 +++--- src/cloudai/parser.py | 17 ++++++++++------ tests/test_cli.py | 9 +++++++++ tests/test_parser.py | 7 +++++++ 5 files changed, 62 insertions(+), 16 deletions(-) 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..d4e34a979 100644 --- a/src/cloudai/parser.py +++ b/src/cloudai/parser.py @@ -43,15 +43,18 @@ 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 +93,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 +111,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..b20be445b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -46,6 +46,15 @@ def test_tests_dir_is_optional(tmp_path: Path): assert "Missing option '--tests-dir'" not in result.output +def test_hook_dir_is_available_on_run_commands(): + runner = CliRunner() + + result = runner.invoke(main, ["dry-run", "--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..028216318 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -39,6 +39,13 @@ 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" + @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" From e34c5b43e3914d457227c6b1f7a9aebe829af20c Mon Sep 17 00:00:00 2001 From: shreyaskommuri Date: Sun, 26 Jul 2026 16:50:35 -0500 Subject: [PATCH 2/3] Address CodeRabbit feedback on hook-dir PR Parametrize test_hook_dir_is_available_on_run_commands over all five common_options commands (install, uninstall, dry-run, run, generate-report) instead of only dry-run, so a regression on any one of them is caught. Add test_custom_hook_root_is_used_by_parse, which exercises Parser.parse() end-to-end with a custom hook_root and asserts a hook test placed only under that directory is actually resolved, rather than just checking the constructor sets hook_root/hook_test_root. Tested: python -m pytest -q (1762 passed, 18 skipped); ruff check; ruff format --check Signed-off-by: shreyaskommuri --- tests/test_cli.py | 5 +++-- tests/test_parser.py | 20 ++++++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index b20be445b..7eda4dc2d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -46,10 +46,11 @@ def test_tests_dir_is_optional(tmp_path: Path): assert "Missing option '--tests-dir'" not in result.output -def test_hook_dir_is_available_on_run_commands(): +@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, ["dry-run", "--help"]) + result = runner.invoke(main, [subcommand, "--help"]) assert result.exit_code == 0 assert "--hook-dir DIRECTORY" in result.output diff --git a/tests/test_parser.py b/tests/test_parser.py index 028216318..aed3c3e20 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -46,6 +46,26 @@ def test_custom_hook_root_is_used(self, parser: Parser, tmp_path: Path): 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} + @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" From 2a85f25608415d14d70617dcaa37e4f77d7c41b5 Mon Sep 17 00:00:00 2001 From: shreyaskommuri Date: Tue, 28 Jul 2026 18:02:41 -0700 Subject: [PATCH 3/3] Address review feedback on hook-dir PR - Remove unused HOOK_TEST_ROOT module constant (superseded by Parser.hook_test_root instance attribute). - Extend custom hook_root test coverage to also exercise scenario-level hook resolution (pre_test/post_test), not just hook test toml loading. --- src/cloudai/parser.py | 1 - tests/test_parser.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/cloudai/parser.py b/src/cloudai/parser.py index d4e34a979..890709d8a 100644 --- a/src/cloudai/parser.py +++ b/src/cloudai/parser.py @@ -37,7 +37,6 @@ from .toml_utils import format_toml_decode_error HOOK_ROOT = Path("conf/hook") -HOOK_TEST_ROOT = HOOK_ROOT / "test" class Parser: diff --git a/tests/test_parser.py b/tests/test_parser.py index aed3c3e20..17e1c11c9 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -66,6 +66,48 @@ def test_custom_hook_root_is_used_by_parse(self, parser: Parser, tmp_path: Path) 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"