Merge pull request #444 from srivatsankrishnan/reporter_bug

Nemo Dry-Run/Run Fix

Author: srivatsankrishnan

src/cloudai/workloads/nemo_run/report_generation_strategy.py

@@ -28,6 +28,10 @@
 
 @cache
 def extract_timings(stdout_file: Path) -> list[float]:
+    if not stdout_file.exists():
+        logging.error(f"{stdout_file} not found")
+        return []
+
     train_step_timings: list[float] = []
     step_timings: list[float] = []
 
@@ -76,6 +80,9 @@ def generate_report(self) -> None:
             return
 
         step_timings = extract_timings(self.results_file)
+        if not step_timings:
+            logging.error("No valid step timings found. Report generation aborted.")
+            return
 
         stats = {
             "avg": np.mean(step_timings),
@@ -93,7 +100,10 @@ def generate_report(self) -> None:
 
     def get_metric(self, metric: str) -> float:
         step_timings = extract_timings(self.results_file)
-        if metric not in {"default", "step-time"} or not step_timings:
+        if not step_timings:
+            return METRIC_ERROR
+
+        if metric not in {"default", "step-time"}:
             return METRIC_ERROR
 
         return float(np.mean(step_timings))

tests/report_generation_strategy/test_nemo_run_report_generation_strategy.py

@@ -22,6 +22,7 @@
 from cloudai import Test, TestRun
 from cloudai.systems.slurm.slurm_system import SlurmSystem
 from cloudai.workloads.nemo_run import NeMoRunCmdArgs, NeMoRunReportGenerationStrategy, NeMoRunTestDefinition
+from cloudai.workloads.nemo_run.report_generation_strategy import extract_timings
 
 
 @pytest.fixture
@@ -154,3 +155,79 @@ def test_metrics(nemo_tr: TestRun, slurm_system: SlurmSystem, metric: str):
     nemo_tr.test.test_definition.agent_metric = metric
     value = nemo_tr.get_metric_value(slurm_system)
     assert value == 12.72090909090909
+
+
+def test_extract_timings_valid_file(tmp_path: Path) -> None:
+    stdout_file = tmp_path / "stdout.txt"
+    stdout_file.write_text(
+        "Training epoch 0, iteration 17/99 | train_step_timing in s: 12.64 | global_step: 17\n"
+        "Training epoch 0, iteration 18/99 | train_step_timing in s: 12.65 | global_step: 18\n"
+        "Training epoch 0, iteration 19/99 | train_step_timing in s: 12.66 | global_step: 19\n"
+    )
+
+    timings = extract_timings(stdout_file)
+    assert timings == [12.65, 12.66], "Timings extraction failed for valid file."
+
+
+def test_extract_timings_missing_file(tmp_path: Path) -> None:
+    stdout_file = tmp_path / "missing_stdout.txt"
+
+    timings = extract_timings(stdout_file)
+    assert timings == [], "Timings extraction should return an empty list for missing file."
+
+
+def test_extract_timings_invalid_content(tmp_path: Path) -> None:
+    stdout_file = tmp_path / "stdout.txt"
+    stdout_file.write_text("Invalid content without timing information\n")
+
+    timings = extract_timings(stdout_file)
+    assert timings == [], "Timings extraction should return an empty list for invalid content."
+
+
+def test_extract_timings_file_not_found(tmp_path: Path) -> None:
+    stdout_file = tmp_path / "nonexistent_stdout.txt"
+
+    timings = extract_timings(stdout_file)
+    assert timings == [], "Timings extraction should return an empty list when the file does not exist."
+
+
+def test_generate_report_no_timings(slurm_system: SlurmSystem, nemo_tr: TestRun, tmp_path: Path) -> None:
+    nemo_tr.output_path = tmp_path
+    stdout_file = nemo_tr.output_path / "stdout.txt"
+    stdout_file.write_text("No valid timing information\n")
+
+    strategy = NeMoRunReportGenerationStrategy(slurm_system, nemo_tr)
+    strategy.generate_report()
+
+    summary_file = nemo_tr.output_path / "report.txt"
+    assert not summary_file.exists(), "Report should not be generated if no valid timings are found."
+
+
+def test_generate_report_partial_timings(slurm_system: SlurmSystem, nemo_tr: TestRun, tmp_path: Path) -> None:
+    nemo_tr.output_path = tmp_path
+    stdout_file = nemo_tr.output_path / "stdout.txt"
+    stdout_file.write_text(
+        "Training epoch 0, iteration 17/99 | train_step_timing in s: 12.64 | global_step: 17\n"
+        "Invalid line without timing\n"
+        "Training epoch 0, iteration 18/99 | train_step_timing in s: 12.65 | global_step: 18\n"
+    )
+
+    strategy = NeMoRunReportGenerationStrategy(slurm_system, nemo_tr)
+    strategy.generate_report()
+
+    summary_file = nemo_tr.output_path / "report.txt"
+    assert summary_file.is_file(), "Report should be generated even with partial valid timings."
+
+    summary_content = summary_file.read_text().strip().split("\n")
+    assert len(summary_content) == 4, "Summary file should contain four lines (avg, median, min, max)."
+
+    expected_values = {
+        "Average": 12.645,
+        "Median": 12.645,
+        "Min": 12.64,
+        "Max": 12.65,
+    }
+
+    for line in summary_content:
+        key, value = line.split(": ")
+        assert pytest.approx(float(value), 0.01) == expected_values[key], f"{key} value mismatch."