Add MegatronBridge post-hook support

src/cloudai/cli/handlers.py

@@ -104,9 +104,7 @@ def prepare_installation(
 ) -> tuple[list[Installable], BaseInstaller]:
     installables: list[Installable] = []
     if scenario:
-        for test in scenario.test_runs:
-            logging.debug(f"{test.test.name} has {len(test.test.installables)} installables.")
-            installables.extend(test.test.installables)
+        installables.extend(_scenario_installables(scenario))
     else:
         for test in tests:
             logging.debug(f"{test.name} has {len(test.installables)} installables.")
@@ -121,6 +119,17 @@ def prepare_installation(
     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()
 

src/cloudai/systems/slurm/slurm_command_gen_strategy.py

@@ -83,7 +83,11 @@ def container_mounts(self) -> list[str]:
 
         repo_mounts = []
         for repo in tdef.git_repos:
-            path = repo.installed_path.absolute() if repo.installed_path else self.system.install_path / repo.repo_name
+            path = (
+                repo.installed_path.absolute()
+                if repo.installed_path
+                else (self.system.install_path / repo.repo_name).absolute()
+            )
             repo_mounts.append(f"{path}:{repo.container_mount}")
 
         mounts = [

src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py

@@ -25,8 +25,10 @@
 
 import toml
 
+from cloudai.core import TestRun, TestScenario
 from cloudai.models.scenario import TestRunDetails
 from cloudai.systems.slurm import SlurmCommandGenStrategy
+from cloudai.util import format_time_limit, parse_time_limit
 
 from .megatron_bridge import HF_TOKEN_REDACTION, MegatronBridgeCmdArgs, MegatronBridgeTestDefinition
 
@@ -88,12 +90,14 @@ def gen_exec_command(self) -> str:
             parts = self._build_launcher_parts(args, tdef, mbridge_repo_path, launcher_py)
 
         launcher_python = str((venv_path / "bin" / "python").absolute())
+        post_hook_sbatch_path = self._gen_post_hook_sbatch() if self.test_run.post_test else None
         full_cmd = self._wrap_launcher_for_job_id_and_quiet_output(
             " ".join(parts),
             launcher_python,
             args.wandb_version,
             args.numpy_version,
             pre_hook_sbatch_path=pre_hook_sbatch_path,
+            post_hook_sbatch_path=post_hook_sbatch_path,
             base_slurm_params=base_slurm_params,
             capture_nodelist=capture_nodelist,
         )
@@ -206,6 +210,105 @@ def _gen_pre_hook_sbatch(self) -> Path:
         sbatch_path.chmod(sbatch_path.stat().st_mode | stat.S_IXUSR)
         return sbatch_path
 
+    def _gen_post_hook_sbatch(self) -> Path:
+        """Generate a standalone sbatch script running post-hook tests."""
+        post_test = self.test_run.post_test
+        if post_test is None:
+            raise RuntimeError("post_test sbatch requested but post_test is not configured.")
+        if not post_test.test_runs:
+            raise RuntimeError("post_test is configured but contains no test runs.")
+
+        post_hook_output = self.test_run.output_path / "post_hook"
+        post_hook_output.mkdir(parents=True, exist_ok=True)
+
+        first_tr = post_test.test_runs[0]
+        first_strategy = self._get_cmd_gen_strategy(first_tr)
+        self._set_hook_output_path(first_tr, self.test_run.output_path / "post_test")
+        first_tr.output_path.mkdir(parents=True, exist_ok=True)
+
+        sbatch_lines = [
+            "#!/bin/bash",
+            f"#SBATCH --job-name=post_hook_{self.job_name()}",
+            f"#SBATCH --output={post_hook_output.absolute() / 'stdout.txt'}",
+            f"#SBATCH --error={post_hook_output.absolute() / 'stderr.txt'}",
+            f"#SBATCH --partition={self.system.default_partition}",
+        ]
+        if self.system.account:
+            sbatch_lines.append(f"#SBATCH --account={self.system.account}")
+        hostfile = self._append_post_hook_resource_directives(first_strategy, post_test, sbatch_lines)
+        if hostfile is not None:
+            sbatch_lines.append(f"export SLURM_HOSTFILE={hostfile}")
+        sbatch_lines.extend(
+            [
+                "",
+                "export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)",
+                "",
+            ]
+        )
+
+        for tr in post_test.test_runs:
+            strategy = first_strategy if tr is first_tr else self._get_cmd_gen_strategy(tr)
+            if tr is not first_tr:
+                self._set_hook_output_path(tr, self.test_run.output_path / "post_test")
+                tr.output_path.mkdir(parents=True, exist_ok=True)
+            srun_command = strategy.gen_srun_command()
+            srun_command_with_output = srun_command.replace(
+                "srun ", f"srun --output={tr.output_path / 'stdout.txt'} --error={tr.output_path / 'stderr.txt'} ", 1
+            )
+            sbatch_lines.append(srun_command_with_output)
+
+        sbatch_path = self.test_run.output_path / "post_hook_sbatch_script.sh"
+        sbatch_path.write_text("\n".join(sbatch_lines))
+        sbatch_path.chmod(sbatch_path.stat().st_mode | stat.S_IXUSR)
+        return sbatch_path
+
+    def _append_post_hook_resource_directives(
+        self,
+        strategy: SlurmCommandGenStrategy,
+        post_test: TestScenario,
+        sbatch_lines: list[str],
+    ) -> Optional[Path]:
+        allocation_run = strategy.test_run
+        original_num_nodes = allocation_run.num_nodes
+        original_nodes = allocation_run.nodes
+        original_exclude_nodes = allocation_run.exclude_nodes
+        allocation_strategy = self._get_cmd_gen_strategy(allocation_run)
+
+        try:
+            allocation_run.num_nodes = self._max_post_hook_nodes(post_test.test_runs)
+            allocation_run.nodes = self._aggregate_post_hook_nodes(post_test.test_runs)
+            allocation_run.exclude_nodes = self._aggregate_post_hook_exclude_nodes(post_test.test_runs)
+            return allocation_strategy._append_resource_directives(
+                sbatch_lines,
+                self._post_hook_time_limit(post_test.test_runs),
+            )
+        finally:
+            allocation_run.num_nodes = original_num_nodes
+            allocation_run.nodes = original_nodes
+            allocation_run.exclude_nodes = original_exclude_nodes
+
+    @staticmethod
+    def _max_post_hook_nodes(test_runs: list[TestRun]) -> int:
+        return max(max(tr.num_nodes) if isinstance(tr.num_nodes, list) else tr.num_nodes for tr in test_runs)
+
+    @staticmethod
+    def _aggregate_post_hook_nodes(test_runs: list[TestRun]) -> list[str]:
+        return list(dict.fromkeys(node for tr in test_runs for node in tr.nodes))
+
+    @staticmethod
+    def _aggregate_post_hook_exclude_nodes(test_runs: list[TestRun]) -> list[str]:
+        return list(dict.fromkeys(node for tr in test_runs for node in tr.exclude_nodes))
+
+    @staticmethod
+    def _post_hook_time_limit(test_runs: list[TestRun]) -> Optional[str]:
+        time_limits = [tr.time_limit for tr in test_runs if tr.time_limit]
+        if not time_limits:
+            return None
+        total_time_limit = parse_time_limit(time_limits[0])
+        for time_limit in time_limits[1:]:
+            total_time_limit += parse_time_limit(time_limit)
+        return format_time_limit(total_time_limit)
+
     def store_test_run(self) -> None:
         test_cmd = self.gen_exec_command()
         trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=test_cmd)
@@ -363,6 +466,7 @@ def _wrap_launcher_for_job_id_and_quiet_output(
         wandb_version: str,
         numpy_version: str,
         pre_hook_sbatch_path: Optional[Path] = None,
+        post_hook_sbatch_path: Optional[Path] = None,
         base_slurm_params: str = "",
         capture_nodelist: bool = False,
     ) -> str:
@@ -374,6 +478,9 @@ def _wrap_launcher_for_job_id_and_quiet_output(
 
         If pre_hook_sbatch_path is provided, the pre-hook sbatch is submitted first and its job ID is used as
         a Slurm dependency (afterok) for the main training job, so training only starts if the pre-hook passed.
+
+        If post_hook_sbatch_path is provided, the post-hook sbatch is submitted with an afterany dependency on
+        the main training job, and CloudAI tracks the post-hook job ID.
         """
         output_dir = self.test_run.output_path.absolute()
         output_dir.mkdir(parents=True, exist_ok=True)
@@ -431,6 +538,20 @@ def _wrap_launcher_for_job_id_and_quiet_output(
         else:
             launch_line = f'{launcher_cmd} >>"$LOG" 2>&1 || LAUNCH_RC=$?'
 
+        post_hook_lines: list[str] = ['  echo "Submitted batch job ${JOB_ID}"']
+        if post_hook_sbatch_path is not None:
+            post_hook_lines = [
+                '  echo "Submitted batch job ${JOB_ID}"',
+                f'  POST_HOOK_SBATCH="{post_hook_sbatch_path.absolute()}"',
+                '  POST_HOOK_OUTPUT=$(sbatch --dependency=afterany:${JOB_ID} "$POST_HOOK_SBATCH" 2>&1)',
+                '  POST_HOOK_JOB_ID=$(echo "$POST_HOOK_OUTPUT" | grep -Eo "Submitted batch job [0-9]+" | grep -Eo "[0-9]+" | tail -n1 || true)',  # noqa: E501
+                '  if [ -z "$POST_HOOK_JOB_ID" ]; then',
+                '    echo "Failed to submit post-hook job: $POST_HOOK_OUTPUT" >&2',
+                "    exit 1",
+                "  fi",
+                '  echo "Submitted post-hook batch job ${POST_HOOK_JOB_ID}"',
+            ]
+
         script_lines = [
             "#!/usr/bin/env bash",
             "set -o pipefail",
@@ -471,7 +592,7 @@ def _wrap_launcher_for_job_id_and_quiet_output(
             '    echo "Megatron-Bridge launcher exited non-zero (${LAUNCH_RC}) after submitting job ${JOB_ID}." >&2',
             '    tail -n 40 "$LOG" >&2 || true',
             "  fi",
-            '  echo "Submitted batch job ${JOB_ID}"',
+            *post_hook_lines,
             "else",
             '  echo "Failed to retrieve job ID." >&2',
             '  if [ "${LAUNCH_RC}" -ne 0 ]; then',

tests/systems/slurm/test_command_gen_strategy.py

@@ -276,6 +276,15 @@ def test_default_container_mounts_with_git_repos(strategy_fixture: SlurmCommandG
     assert mounts[4] == f"{repo2.installed_path}:{repo2.container_mount}"
 
 
+def test_default_container_mounts_with_uninstalled_git_repo(strategy_fixture: SlurmCommandGenStrategy):
+    repo = GitRepo(url="./git_repo", commit="commit", mount_as="/git/repo")
+    strategy_fixture.test_run.test.git_repos = [repo]
+
+    mounts = strategy_fixture.container_mounts()
+
+    assert mounts[3] == f"{(strategy_fixture.system.install_path / repo.repo_name).absolute()}:{repo.container_mount}"
+
+
 def test_ranks_mapping_cmd(strategy_fixture: SlurmCommandGenStrategy):
     expected_command = (
         f"srun --export=ALL --mpi={strategy_fixture.system.mpi} -N{strategy_fixture.test_run.num_nodes} "

tests/test_handlers.py

@@ -26,6 +26,7 @@
 
 from cloudai.cli.handlers import (
     handle_dse_job,
+    prepare_installation,
     validate_domain_randomization_active,
     verify_system_configs,
     verify_test_configs,
@@ -35,6 +36,7 @@
 from cloudai.core import (
     BaseAgent,
     BaseAgentConfig,
+    GitRepo,
     Parser,
     Registry,
     RewardOverrides,
@@ -45,6 +47,7 @@
     TestScenarioParsingError,
 )
 from cloudai.models.scenario import ReportConfig
+from cloudai.models.workload import CmdArgs, TestDefinition
 from cloudai.reporter import StatusReporter
 from cloudai.systems.slurm.slurm_system import SlurmSystem
 
@@ -166,6 +169,64 @@ def test_dse_run_uses_agent_config(
     assert recorded.random_seed == expected["random_seed"]
 
 
+def test_prepare_installation_includes_hook_installables(slurm_system: SlurmSystem) -> None:
+    parent_repo = GitRepo(url="./parent", commit="main")
+    pre_repo = GitRepo(url="./pre", commit="main")
+    post_repo = GitRepo(url="./post", commit="main")
+    parent_run = TestRun(
+        name="parent",
+        test=TestDefinition(
+            name="parent",
+            description="parent",
+            test_template_name="template",
+            cmd_args=CmdArgs(),
+            git_repos=[parent_repo],
+        ),
+        num_nodes=1,
+        nodes=[],
+        pre_test=TestScenario(
+            name="pre",
+            test_runs=[
+                TestRun(
+                    name="pre",
+                    test=TestDefinition(
+                        name="pre",
+                        description="pre",
+                        test_template_name="template",
+                        cmd_args=CmdArgs(),
+                        git_repos=[pre_repo],
+                    ),
+                    num_nodes=1,
+                    nodes=[],
+                )
+            ],
+        ),
+        post_test=TestScenario(
+            name="post",
+            test_runs=[
+                TestRun(
+                    name="post",
+                    test=TestDefinition(
+                        name="post",
+                        description="post",
+                        test_template_name="template",
+                        cmd_args=CmdArgs(),
+                        git_repos=[post_repo],
+                    ),
+                    num_nodes=1,
+                    nodes=[],
+                )
+            ],
+        ),
+    )
+
+    installables, _ = prepare_installation(slurm_system, [], TestScenario(name="scenario", test_runs=[parent_run]))
+
+    assert parent_repo in installables
+    assert pre_repo in installables
+    assert post_repo in installables
+
+
 def test_dse_run_cache(base_tr: TestRun, tmp_path, caplog: pytest.LogCaptureFixture):
     base_tr.test.cmd_args.candidate = [1, 1, 2]
     base_tr.test.agent = "grid_search"