diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index a3daa9f3e..83bb493d2 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -17,6 +17,7 @@ import logging import os import re +from pathlib import Path from typing import List, Optional, Union, cast from pydantic import Field, ValidationInfo, field_validator @@ -581,11 +582,11 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: - At the point of failure the script asks for reference golden values, that we don't have - Then the script will perform convergence test between provided golden and actual golden - we don't need it """ - log_path = tr.output_path / "cloudai_megatron_bridge_launcher.log" - if not log_path.is_file(): + log_path = find_mbridge_log(tr.output_path) + if log_path is None: return JobStatusResult( is_successful=False, - error_message=f"Megatron-Bridge launcher log not found in {tr.output_path}.", + error_message=f"Megatron-Bridge training log not found in {tr.output_path}.", ) log_data = log_path.read_text(encoding="utf-8", errors="ignore") @@ -596,6 +597,15 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: return JobStatusResult(is_successful=True) +def find_mbridge_log(output_path: Path) -> Path | None: + training_logs = sorted(output_path.glob("experiments/**/log*.out")) + if training_logs: + return training_logs[-1] + + launcher_log = output_path / "cloudai_megatron_bridge_launcher.log" + return launcher_log if launcher_log.is_file() else None + + def extract_mbridge_metrics(logs: str) -> tuple[list[float], list[float]]: step_times_s: list[float] = [] gpu_tflops: list[float] = [] diff --git a/src/cloudai/workloads/megatron_bridge/report_generation_strategy.py b/src/cloudai/workloads/megatron_bridge/report_generation_strategy.py index c21e170d3..52337d330 100644 --- a/src/cloudai/workloads/megatron_bridge/report_generation_strategy.py +++ b/src/cloudai/workloads/megatron_bridge/report_generation_strategy.py @@ -20,7 +20,7 @@ from cloudai.core import METRIC_ERROR, MetricValue, ReportGenerationStrategy -from .megatron_bridge import extract_mbridge_metrics +from .megatron_bridge import extract_mbridge_metrics, find_mbridge_log class MegatronBridgeReportGenerationStrategy(ReportGenerationStrategy): @@ -29,8 +29,7 @@ class MegatronBridgeReportGenerationStrategy(ReportGenerationStrategy): metrics: ClassVar[list[str]] = ["default", "step-time", "tflops-per-gpu"] def get_log_file(self) -> Path | None: - log = self.test_run.output_path / "cloudai_megatron_bridge_launcher.log" - return log if log.is_file() else None + return find_mbridge_log(self.test_run.output_path) @property def results_file(self) -> Path: @@ -52,7 +51,7 @@ def generate_report(self) -> None: log_file, step_times_s, gpu_tflops = self._get_extracted_data() if not log_file: logging.error( - "No Megatron-Bridge launcher log file found in: %s", + "No Megatron-Bridge training log file found in: %s", self.test_run.output_path, ) return @@ -107,7 +106,7 @@ def get_metric(self, metric: str) -> MetricValue: log_file, step_times_s, gpu_tflops = self._get_extracted_data() if not log_file: logging.error( - "No Megatron-Bridge launcher log file found in: %s", + "No Megatron-Bridge training log file found in: %s", self.test_run.output_path, ) return METRIC_ERROR 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 31cf500f4..570eac767 100644 --- a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py @@ -17,8 +17,10 @@ from __future__ import annotations import logging +import os import shlex import stat +import subprocess from pathlib import Path from typing import Any, Optional, cast @@ -50,6 +52,10 @@ def _container_mounts(self) -> list[str]: return [] def gen_exec_command(self) -> str: + existing_script = next(self.test_run.output_path.glob("experiments/**/*_sbatch.sh"), None) + if existing_script: + return f"sbatch {shlex.quote(str(existing_script))}" + tdef: MegatronBridgeTestDefinition = cast(MegatronBridgeTestDefinition, self.test_run.test) args: MegatronBridgeCmdArgs = tdef.cmd_args @@ -73,32 +79,61 @@ def gen_exec_command(self) -> str: launcher_py = (mbridge_repo_path / "scripts" / "performance" / "setup_experiment.py").absolute() - 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, - args.wandb_version, - args.numpy_version, - pre_hook_sbatch_path=pre_hook_sbatch_path, - base_slurm_params=base_slurm_params, - capture_nodelist=capture_nodelist, - ) + patcher_path = self._write_dryrun_patcher(launcher_py) + parts = self._build_launcher_parts(args, tdef, mbridge_repo_path, patcher_path) + if "-d" not in parts: + parts.append("-d") + launcher_cmd = " ".join(parts) + log_path = self.test_run.output_path / "cloudai_megatron_bridge_launcher.log" + env = os.environ.copy() + for key in self.CONTAINER_RUNTIME_ENV_VARS: + if key in self.final_env_vars: + env[key] = str(self.final_env_vars[key]) + with log_path.open("w") as log: + result = subprocess.run( + launcher_cmd, + shell=True, + executable="/bin/bash", + env=env, + stdout=log, + stderr=subprocess.STDOUT, + ) + if result.returncode: + raise RuntimeError(f"Megatron-Bridge dry-run failed. See {log_path}.") + + script = next(self.test_run.output_path.glob("experiments/**/*_sbatch.sh"), None) + if script is None: + raise RuntimeError(f"Megatron-Bridge dry-run did not generate an sbatch script. See {log_path}.") + full_cmd = f"sbatch {shlex.quote(str(script))}" self._write_command_to_file(full_cmd, self.test_run.output_path) return full_cmd + def _write_dryrun_patcher(self, launcher_py: Path) -> Path: + patcher_path = self.test_run.output_path / "cloudai_megatron_bridge_dryrun.py" + patcher_path.write_text( + "\n".join( + [ + "import runpy", + "import sys", + "from pathlib import Path", + "from nemo_run.run.experiment import Experiment", + "", + "original_dryrun = Experiment.dryrun", + "", + "def keep_dryrun_artifacts(self, log=True, exist_ok=False, delete_exp_dir=True):", + " return original_dryrun(self, log=log, exist_ok=exist_ok, delete_exp_dir=False)", + "", + "Experiment.dryrun = keep_dryrun_artifacts", + f"launcher = Path({str(launcher_py)!r})", + "sys.path.insert(0, str(launcher.parent))", + 'runpy.run_path(str(launcher), run_name="__main__")', + "", + ] + ) + ) + return patcher_path + def _collect_additional_slurm_params(self) -> list[str]: """Return the additional_slurm_params list (without dependency).""" params: list[str] = []