Allow external callers to configure the hook directory

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,15 +237,19 @@ 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.
 
     While this process is automatically executed as part of "run" command, one can also invoke it manually using this
     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))
 

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
 

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

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"]
 )

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 `<hook_root>/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"