[MBridge] CloudAI owns sbatch submission

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] = []

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

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] = []