diff --git a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py index 7270cbac7..61dfdda01 100644 --- a/src/cloudai/systems/slurm/slurm_command_gen_strategy.py +++ b/src/cloudai/systems/slurm/slurm_command_gen_strategy.py @@ -372,6 +372,41 @@ def _write_sbatch_script(self, srun_command: str) -> str: return f"sbatch {batch_script_path}" + def _append_resource_directives(self, content: List[str], time_limit: Optional[str] = None) -> Optional[Path]: + """ + Append resource-related SBATCH directives shared between main and pre-hook scripts. + + Covers: reservation, distribution, nodelist/exclude, gpus-per-node/gres, + ntasks-per-node, time limit, and extra_sbatch_args. Does NOT write + job-name, output/error paths, partition, or account — those differ + between the main job and the pre-hook and are the caller's responsibility. + + Args: + content: The list of script lines to append to. + time_limit: SBATCH time limit string; uses self.test_run.time_limit if None. + + Returns: + Path to the generated hostfile, or None if node-list mode was not used. + """ + self._add_reservation(content) + hostfile = self._append_nodes_related_directives(content) + + if self.system.gpus_per_node and self.system.supports_gpu_directives: + content.append(f"#SBATCH --gpus-per-node={self.system.gpus_per_node}") + content.append(f"#SBATCH --gres=gpu:{self.system.gpus_per_node}") + + if self.system.ntasks_per_node: + content.append(f"#SBATCH --ntasks-per-node={self.system.ntasks_per_node}") + + limit = time_limit if time_limit is not None else self.test_run.time_limit + if limit: + content.append(f"#SBATCH --time={limit}") + + for arg in self.system.extra_sbatch_args: + content.append(f"#SBATCH {arg}") + + return hostfile + def _append_sbatch_directives(self, batch_script_content: List[str]) -> None: """ Append SBATCH directives to the batch script content. @@ -379,27 +414,13 @@ def _append_sbatch_directives(self, batch_script_content: List[str]) -> None: Args: batch_script_content (List[str]): The list of script lines to append to. """ - batch_script_content = self._add_reservation(batch_script_content) - batch_script_content.append(f"#SBATCH --output={self.test_run.output_path.absolute() / 'stdout.txt'}") batch_script_content.append(f"#SBATCH --error={self.test_run.output_path.absolute() / 'stderr.txt'}") batch_script_content.append(f"#SBATCH --partition={self.system.default_partition}") if self.system.account: batch_script_content.append(f"#SBATCH --account={self.system.account}") - hostfile = self._append_nodes_related_directives(batch_script_content) - - if self.system.gpus_per_node and self.system.supports_gpu_directives: - batch_script_content.append(f"#SBATCH --gpus-per-node={self.system.gpus_per_node}") - batch_script_content.append(f"#SBATCH --gres=gpu:{self.system.gpus_per_node}") - - if self.system.ntasks_per_node: - batch_script_content.append(f"#SBATCH --ntasks-per-node={self.system.ntasks_per_node}") - if self.test_run.time_limit: - batch_script_content.append(f"#SBATCH --time={self.test_run.time_limit}") - - for arg in self.system.extra_sbatch_args: - batch_script_content.append(f"#SBATCH {arg}") + hostfile = self._append_resource_directives(batch_script_content) if hostfile is not None: batch_script_content.append(f"export SLURM_HOSTFILE={hostfile}") 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 d07e4bf77..d0f24d3e9 100644 --- a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py @@ -73,13 +73,135 @@ def gen_exec_command(self) -> str: launcher_py = (mbridge_repo_path / "scripts" / "performance" / "setup_experiment.py").absolute() - parts = self._build_launcher_parts(args, tdef, mbridge_repo_path, launcher_py) + pre_hook_sbatch_path: Optional[Path] = None + base_slurm_params: str = "" + capture_nodelist: bool = False + if self.test_run.pre_test: + pre_hook_sbatch_path = self._gen_pre_hook_sbatch() + parts = self._build_launcher_parts(args, tdef, mbridge_repo_path, launcher_py, include_slurm_params=False) + base_slurm_params = ";".join(self._collect_additional_slurm_params()) + _, node_list = self.get_cached_nodes_spec() + capture_nodelist = not node_list + else: + parts = self._build_launcher_parts(args, tdef, mbridge_repo_path, launcher_py) + launcher_python = str((venv_path / "bin" / "python").absolute()) - full_cmd = self._wrap_launcher_for_job_id_and_quiet_output(" ".join(parts), launcher_python) + full_cmd = self._wrap_launcher_for_job_id_and_quiet_output( + " ".join(parts), + launcher_python, + pre_hook_sbatch_path=pre_hook_sbatch_path, + base_slurm_params=base_slurm_params, + capture_nodelist=capture_nodelist, + ) self._write_command_to_file(full_cmd, self.test_run.output_path) return full_cmd + def _collect_additional_slurm_params(self) -> list[str]: + """Return the additional_slurm_params list (without dependency).""" + params: list[str] = [] + if self.system.gpus_per_node and self.system.supports_gpu_directives: + params.append(f"gpus-per-node={self.system.gpus_per_node}") + params.append(f"gres=gpu:{self.system.gpus_per_node}") + _, node_list = self.get_cached_nodes_spec() + if node_list: + params.append(f"nodelist={','.join(node_list)}") + elif self.test_run.exclude_nodes: + params.append(f"exclude={','.join(self.test_run.exclude_nodes)}") + for source in (self.system.extra_srun_args, self.test_run.extra_srun_args): + if source: + params.extend(self._parse_srun_args_as_slurm_params(source)) + return params + + def _gen_pre_hook_sbatch(self) -> Path: + """ + Generate a standalone sbatch script running per-node independent pre-hook tests. + + Each node runs its own alltoall among its local GPUs (1 srun per node in parallel), + so the tests are truly independent — no cross-node NCCL communicator is formed. + """ + pre_hook_output = self.test_run.output_path / "pre_hook" + pre_hook_output.mkdir(parents=True, exist_ok=True) + + sbatch_lines = [ + "#!/bin/bash", + f"#SBATCH --job-name=pre_hook_{self.job_name()}", + f"#SBATCH --output={pre_hook_output.absolute() / 'stdout.txt'}", + f"#SBATCH --error={pre_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}") + self._append_resource_directives(sbatch_lines, self.test_run.time_limit) + if self.test_run.extra_srun_args: + for param in self._parse_srun_args_as_slurm_params(self.test_run.extra_srun_args): + key, _, val = param.partition("=") + sbatch_lines.append(f"#SBATCH --{key}={val}" if val else f"#SBATCH --{key}") + sbatch_lines.append("") + + assert self.test_run.pre_test is not None + success_vars = [] + for idx, tr in enumerate(self.test_run.pre_test.test_runs): + tr.num_nodes = 1 + strategy = self._get_cmd_gen_strategy(tr) + self._set_hook_output_path(tr, self.test_run.output_path / "pre_test") + tr.output_path.mkdir(parents=True, exist_ok=True) + + node_out = tr.output_path.absolute() + srun_cmd = strategy.gen_srun_command() + # Inject per-node output paths and --nodelist so each node runs independently + node_srun = srun_cmd.replace( + "srun ", + f"srun --nodelist=$_node --output={node_out}/stdout_$_node.txt --error={node_out}/stderr_$_node.txt ", + 1, + ) + + success_var = f"SUCCESS_{idx}" + success_vars.append(success_var) + + min_busbw = getattr(tr.test.cmd_args, "min_busbw", None) + if min_busbw is not None: + check_cmd = ( + f"awk '/Avg bus bandwidth/ {{ if ($NF+0 >= {min_busbw}) found=1 }}" + f" END {{ exit !found }}' {node_out}/stdout_$_node.txt 2>/dev/null" + ) + else: + check_cmd = f'grep -q "Avg bus bandwidth" {node_out}/stdout_$_node.txt 2>/dev/null' + + sbatch_lines.extend( + [ + f"# {tr.test.name}", + f"mkdir -p {node_out}", + "for _node in $(scontrol show hostnames $SLURM_JOB_NODELIST); do", + f" {node_srun} &", + "done", + "wait", + f"{success_var}=1", + "for _node in $(scontrol show hostnames $SLURM_JOB_NODELIST); do", + f" if ! {check_cmd}; then", + f" {success_var}=0", + " fi", + "done", + "", + ] + ) + + combined = " && ".join([f"[ ${v} -eq 1 ]" for v in success_vars]) + sbatch_lines.extend( + [ + f"PRE_TEST_SUCCESS=$( {combined} && echo 1 || echo 0 )", + 'if [ "$PRE_TEST_SUCCESS" -ne 1 ]; then', + ' echo "Pre-hook tests failed. Blocking training job." >&2', + " exit 1", + "fi", + ] + ) + + sbatch_path = self.test_run.output_path / "pre_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 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) @@ -166,12 +288,22 @@ def _normalize_cuda_graph_scope_arg(self, val: Any) -> str: parts = [p.strip().strip("\"'") for p in s.split(",") if p.strip()] return ",".join(parts) - def _wrap_launcher_for_job_id_and_quiet_output(self, launcher_cmd: str, launcher_python: str) -> str: + def _wrap_launcher_for_job_id_and_quiet_output( + self, + launcher_cmd: str, + launcher_python: str, + pre_hook_sbatch_path: Optional[Path] = None, + base_slurm_params: str = "", + capture_nodelist: bool = False, + ) -> str: """ Run the Megatron-Bridge launcher quietly and ensure CloudAI can parse a job ID. CloudAI's SlurmRunner expects stdout to include "Submitted batch job ". This writes a readable wrapper script (with section breaks) into the test output directory, then runs it. + + 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. """ output_dir = self.test_run.output_path.absolute() output_dir.mkdir(parents=True, exist_ok=True) @@ -181,6 +313,54 @@ def _wrap_launcher_for_job_id_and_quiet_output(self, launcher_cmd: str, launcher container_runtime_exports = self._container_runtime_env_exports() + pre_hook_lines: list[str] = [] + launch_line: str + if pre_hook_sbatch_path is not None: + nodelist_lines: list[str] = [] + if capture_nodelist: + nodelist_lines = [ + "# Wait for pre-hook to reach RUNNING; exit on terminal state or timeout (1 h)", + "PRE_HOOK_NODES=''", + "_NODELIST_WAIT=0", + "_NODELIST_TIMEOUT=3600", + "while true; do", + ' _state=$(squeue -j "$PRE_HOOK_JOB_ID" -h -o "%T" 2>/dev/null || true)', + ' if [ "$_state" = "RUNNING" ]; then', + ' PRE_HOOK_NODES=$(squeue -j "$PRE_HOOK_JOB_ID" -h -o "%N" 2>/dev/null | head -1 || true)', + " break", + ' elif [ -z "$_state" ] || [ "$_state" = "FAILED" ] || [ "$_state" = "CANCELLED" ] || [ "$_state" = "COMPLETED" ] || [ "$_state" = "TIMEOUT" ]; then', # noqa: E501 + " echo \"Pre-hook job $PRE_HOOK_JOB_ID ended in state '${_state:-gone}' before reaching RUNNING.\" >&2", # noqa: E501 + " exit 1", + ' elif [ "$_NODELIST_WAIT" -ge "$_NODELIST_TIMEOUT" ]; then', + ' echo "Timed out after ${_NODELIST_TIMEOUT}s waiting for pre-hook job $PRE_HOOK_JOB_ID to reach RUNNING (last state: ${_state})." >&2', # noqa: E501 + " exit 1", + " fi", + " sleep 10", + " _NODELIST_WAIT=$((_NODELIST_WAIT + 10))", + "done", + 'ADDITIONAL_SLURM_PARAMS="${ADDITIONAL_SLURM_PARAMS};nodelist=${PRE_HOOK_NODES}"', + "", + ] + pre_hook_lines = [ + f'PRE_HOOK_SBATCH="{pre_hook_sbatch_path.absolute()}"', + 'PRE_HOOK_OUTPUT=$(sbatch "$PRE_HOOK_SBATCH" 2>&1)', + 'PRE_HOOK_JOB_ID=$(echo "$PRE_HOOK_OUTPUT" | grep -Eo "Submitted batch job [0-9]+" | grep -Eo "[0-9]+" | tail -n1 || true)', # noqa: E501 + 'if [ -z "$PRE_HOOK_JOB_ID" ]; then', + ' echo "Failed to submit pre-hook job: $PRE_HOOK_OUTPUT" >&2', + " exit 1", + "fi", + 'echo "Submitted pre-hook batch job $PRE_HOOK_JOB_ID"', + f'ADDITIONAL_SLURM_PARAMS="{base_slurm_params}"', + 'ADDITIONAL_SLURM_PARAMS="${ADDITIONAL_SLURM_PARAMS};dependency=afterok:${PRE_HOOK_JOB_ID}"', + "", + *nodelist_lines, + ] + launch_line = ( + f'{launcher_cmd} --additional_slurm_params "$ADDITIONAL_SLURM_PARAMS" >>"$LOG" 2>&1 || LAUNCH_RC=$?' + ) + else: + launch_line = f'{launcher_cmd} >>"$LOG" 2>&1 || LAUNCH_RC=$?' + script_lines = [ "#!/usr/bin/env bash", "set -o pipefail", @@ -195,6 +375,7 @@ def _wrap_launcher_for_job_id_and_quiet_output(self, launcher_cmd: str, launcher "", *container_runtime_exports, "", + *pre_hook_lines, ': >"$LOG"', "WANDB_INSTALL_RC=0", f'{shlex.quote(launcher_python)} -m pip install wandb numpy==1.26.4 >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$?', @@ -205,7 +386,7 @@ def _wrap_launcher_for_job_id_and_quiet_output(self, launcher_cmd: str, launcher "fi", "", "LAUNCH_RC=0", - f'{launcher_cmd} >>"$LOG" 2>&1 || LAUNCH_RC=$?', + launch_line, "", # Parse job id from Megatron-Bridge output (multiple possible formats) # Patterns: "Submitted batch job 694112", "Job id: 694112", "- Job id: 694112", "Job ID: 694112" @@ -257,7 +438,12 @@ def _add_extra_cmd_args(self, extra_cmd_args: dict[str, str]) -> list[str]: return [shlex.quote(f"{key}={value}" if value else key) for key, value in overrides.items()] def _build_launcher_parts( # noqa: C901 - self, args: MegatronBridgeCmdArgs, tdef: MegatronBridgeTestDefinition, repo_path: Path, launcher_py: Path + self, + args: MegatronBridgeCmdArgs, + tdef: MegatronBridgeTestDefinition, + repo_path: Path, + launcher_py: Path, + include_slurm_params: bool = True, ) -> list[str]: fields_set = args.model_fields_set force_fields = { @@ -462,25 +648,10 @@ def add_field(field: str, flag: str, value: Any) -> None: add_field("nsys_trace", "--nsys_trace", self._list_or_comma_str(args.nsys_trace)) add_field("nsys_extra_args", "--nsys_extra_args", self._list_or_comma_str(args.nsys_extra_args)) - additional_slurm_params: list[str] = [] - - if self.system.gpus_per_node and self.system.supports_gpu_directives: - additional_slurm_params.append(f"gpus-per-node={self.system.gpus_per_node}") - additional_slurm_params.append(f"gres=gpu:{self.system.gpus_per_node}") - - _, node_list = self.get_cached_nodes_spec() - if node_list: - nodelist_str = ",".join(node_list) - additional_slurm_params.append(f"nodelist={nodelist_str}") - elif self.test_run.exclude_nodes: - additional_slurm_params.append(f"exclude={','.join(self.test_run.exclude_nodes)}") - - for source in (self.system.extra_srun_args, self.test_run.extra_srun_args): - if source: - additional_slurm_params.extend(self._parse_srun_args_as_slurm_params(source)) - - if additional_slurm_params: - parts.extend(["--additional_slurm_params", shlex.quote(";".join(additional_slurm_params))]) + if include_slurm_params: + additional_slurm_params = self._collect_additional_slurm_params() + if additional_slurm_params: + parts.extend(["--additional_slurm_params", shlex.quote(";".join(additional_slurm_params))]) # Config variant add_field("config_variant", "-cv", args.config_variant) diff --git a/src/cloudai/workloads/nccl_test/nccl.py b/src/cloudai/workloads/nccl_test/nccl.py index e54d86c7d..7308c5a56 100644 --- a/src/cloudai/workloads/nccl_test/nccl.py +++ b/src/cloudai/workloads/nccl_test/nccl.py @@ -103,6 +103,7 @@ class NCCLCmdArgs(CmdArgs): stepfactor: Optional[Union[int, list[int]]] = None use_deepep_matrix: bool = False alltoallv_matrix_container_path: str = "/tmp/traffic_matrix.txt" + min_busbw: Optional[float] = None class NCCLTestDefinition(TestDefinition): diff --git a/src/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py b/src/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py index ce1fffeb1..d48fcff2a 100644 --- a/src/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nccl_test/slurm_command_gen_strategy.py @@ -31,6 +31,7 @@ "subtest_name", "use_deepep_matrix", "alltoallv_matrix_container_path", + "min_busbw", } @@ -148,11 +149,26 @@ def _alltoallv_test_command(self, tdef: NCCLTestDefinition) -> List[str]: parts.append(self.test_run.test.extra_args_str) return parts + @property + def _is_single_process(self) -> bool: + tdef: NCCLTestDefinition = cast(NCCLTestDefinition, self.test_run.test) + name = tdef.cmd_args.subtest_name + if isinstance(name, list): + name = name[0] + return not name.endswith("_mpi") + + def gen_srun_prefix(self, use_pretest_extras: bool = False, with_num_nodes: bool = True) -> List[str]: + parts = super().gen_srun_prefix(use_pretest_extras=use_pretest_extras, with_num_nodes=with_num_nodes) + if self._is_single_process: + parts = ["--ntasks=1" if p == f"--mpi={self.mpi}" else p for p in parts] + return parts + def generate_test_command(self) -> List[str]: tdef: NCCLTestDefinition = cast(NCCLTestDefinition, self.test_run.test) if tdef.cmd_args.subtest_name == "alltoallv_perf_mpi": return self._alltoallv_test_command(tdef) + single_process = self._is_single_process srun_command_parts = [f"{tdef.cmd_args.subtest_name}"] nccl_test_args = tdef.cmd_args.model_dump().keys() for arg in nccl_test_args: @@ -163,7 +179,9 @@ def generate_test_command(self) -> List[str]: if value is None: continue - if len(arg) > 1: + if arg == "ngpus" and single_process: + srun_command_parts.append(f"-g {value}") + elif len(arg) > 1: srun_command_parts.append(f"--{arg} {value}") else: srun_command_parts.append(f"-{arg} {value}")