diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index b3ed53071..1fc620544 100644 --- a/src/cloudai/cli/handlers.py +++ b/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() diff --git a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py index 2c5067954..be4b57a3c 100644 --- a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py +++ b/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 = [ diff --git a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py b/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py index 40196824c..280b61242 100644 --- a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py +++ b/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', diff --git a/tests/systems/slurm/test_command_gen_strategy.py b/tests/systems/slurm/test_command_gen_strategy.py index 98e103704..642b63929 100644 --- a/tests/systems/slurm/test_command_gen_strategy.py +++ b/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} " diff --git a/tests/test_handlers.py b/tests/test_handlers.py index bd2a76c46..560b2f409 100644 --- a/tests/test_handlers.py +++ b/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" diff --git a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py index d3d1e1e7b..55dda104b 100644 --- a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py +++ b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py @@ -20,7 +20,7 @@ import pytest import toml -from cloudai.core import GitRepo, TestRun +from cloudai.core import GitRepo, TestRun, TestScenario from cloudai.systems.slurm import SlurmSystem from cloudai.workloads.megatron_bridge import ( MegatronBridgeCmdArgs, @@ -28,6 +28,7 @@ MegatronBridgeTestDefinition, ) from cloudai.workloads.megatron_bridge.megatron_bridge import HF_TOKEN_REDACTION +from cloudai.workloads.nccl_test import NCCLCmdArgs, NCCLTestDefinition WRAPPER_SCRIPT_NAME = "cloudai_megatron_bridge_submit_and_parse_jobid.sh" @@ -338,6 +339,117 @@ def test_wrapper_emits_job_id_even_when_launcher_non_zero( assert 'exit "${LAUNCH_RC}"' not in wrapper_content assert "Submitted batch job[ ]+[0-9]+" in wrapper_content + def test_post_hook_runs_as_dependent_job( + self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun], tmp_path: Path + ) -> None: + tr = make_test_run(output_subdir="out_post_hook") + hook_tdef = NCCLTestDefinition( + name="nccl_post", + description="post", + test_template_name="NcclTest", + cmd_args=NCCLCmdArgs(docker_image_url="fake://url/nccl"), + extra_env_vars={"HOOK_VAR": "1"}, + ) + post_run = TestRun( + test=hook_tdef, + name="nccl_post", + num_nodes=1, + nodes=[], + output_path=tmp_path / "unused", + time_limit="00:05:00", + ) + tr.post_test = TestScenario(name="post", test_runs=[post_run]) + + cmd_gen = MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr) + wrapper_content = self._wrapper_content(cmd_gen) + post_hook_script = tr.output_path / "post_hook_sbatch_script.sh" + + assert post_hook_script.exists() + post_hook_content = post_hook_script.read_text() + assert "#SBATCH --time=00:05:00" in post_hook_content + assert "/post_test/nccl_post/stdout.txt" in post_hook_content + assert "srun " in post_hook_content + assert "POST_HOOK_OUTPUT=$(sbatch --dependency=afterany:${JOB_ID}" in wrapper_content + assert 'echo "Submitted post-hook batch job ${POST_HOOK_JOB_ID}"' in wrapper_content + assert 'echo "Submitted batch job ${JOB_ID}"' in wrapper_content + assert 'echo "Submitted batch job ${POST_HOOK_JOB_ID}"' not in wrapper_content + + def test_post_hook_uses_largest_allocation( + self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun], tmp_path: Path + ) -> None: + tr = make_test_run(output_subdir="out_post_hook_resources") + first_post = TestRun( + test=NCCLTestDefinition( + name="post_one", + description="post", + test_template_name="NcclTest", + cmd_args=NCCLCmdArgs(docker_image_url="fake://url/nccl"), + ), + name="post_one", + num_nodes=1, + nodes=[], + output_path=tmp_path / "unused_one", + time_limit="00:05:00", + ) + second_post = TestRun( + test=NCCLTestDefinition( + name="post_two", + description="post", + test_template_name="NcclTest", + cmd_args=NCCLCmdArgs(docker_image_url="fake://url/nccl"), + ), + name="post_two", + num_nodes=3, + nodes=[], + output_path=tmp_path / "unused_two", + time_limit="00:10:00", + ) + tr.post_test = TestScenario(name="post", test_runs=[first_post, second_post]) + + self._wrapper_content(MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr)) + post_hook_content = (tr.output_path / "post_hook_sbatch_script.sh").read_text() + + assert "#SBATCH -N 3" in post_hook_content + assert "#SBATCH --time=00:15:00" in post_hook_content + assert "/post_test/post_one/stdout.txt" in post_hook_content + assert "/post_test/post_two/stdout.txt" in post_hook_content + post_one_srun = next( + line for line in post_hook_content.splitlines() if "/post_test/post_one/stdout.txt" in line + ) + post_two_srun = next( + line for line in post_hook_content.splitlines() if "/post_test/post_two/stdout.txt" in line + ) + assert " -N1 " in post_one_srun + assert " -N3 " in post_two_srun + + def test_post_hook_exports_hostfile_for_explicit_nodes( + self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun], tmp_path: Path + ) -> None: + configured_slurm_system.ntasks_per_node = 2 + tr = make_test_run(output_subdir="out_post_hook_hostfile") + post_run = TestRun( + test=NCCLTestDefinition( + name="nccl_post", + description="post", + test_template_name="NcclTest", + cmd_args=NCCLCmdArgs(docker_image_url="fake://url/nccl"), + ), + name="nccl_post", + num_nodes=1, + nodes=["node2", "node1"], + output_path=tmp_path / "unused", + time_limit="00:05:00", + ) + tr.post_test = TestScenario(name="post", test_runs=[post_run]) + + self._wrapper_content(MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr)) + post_hook_content = (tr.output_path / "post_hook_sbatch_script.sh").read_text() + hostfile_path = tr.output_path / "post_test" / "nccl_post" / "hostfile.txt" + + assert "#SBATCH --nodelist=node1,node2" in post_hook_content + assert f"export SLURM_HOSTFILE={hostfile_path.absolute()}" in post_hook_content + assert hostfile_path.read_text().splitlines() == ["node1", "node1", "node2", "node2"] + def test_wrapper_installs_wandb_before_launcher( self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun] ) -> None: