feat(aorta): collect per-node traces, survive partial node failure

cvs/parsers/aorta_report.py

@@ -76,24 +76,35 @@ def parse(self, run_result: RunResult) -> ParseResult[AortaTraceMetrics]:
         Returns:
             ParseResult containing validated AortaTraceMetrics for each rank
         """
+        run_warnings = []
         if not run_result.succeeded:
-            return ParseResult(status=ParseStatus.FAILED, errors=[f"Run did not succeed: {run_result.error_message}"])
+            # A failed/timed-out node no longer discards traces collected from
+            # surviving nodes (see AortaRunner.run()), so a partial run can
+            # still have real reports to parse. Only bail out below if there
+            # is nothing on disk to parse.
+            run_warnings.append(f"Run did not succeed: {run_result.error_message}")
+            log.warning(run_warnings[-1])
 
         # First try the tracelens_analysis artifact
         analysis_dir = run_result.get_artifact("tracelens_analysis")
         if analysis_dir and analysis_dir.exists():
-            return self.parse_analysis_directory(analysis_dir)
+            result = self.parse_analysis_directory(analysis_dir)
+            result.warnings = run_warnings + list(result.warnings)
+            return result
 
         # Fallback: look for analysis dir relative to trace dir
         trace_dir = run_result.get_artifact("torch_traces")
         if trace_dir:
             parent_dir = trace_dir.parent
             analysis_dir = parent_dir / "tracelens_analysis"
             if analysis_dir.exists():
-                return self.parse_analysis_directory(analysis_dir)
+                result = self.parse_analysis_directory(analysis_dir)
+                result.warnings = run_warnings + list(result.warnings)
+                return result
 
         return ParseResult(
             status=ParseStatus.FAILED,
+            warnings=run_warnings,
             errors=["No tracelens_analysis artifact found. Ensure analysis.enable_tracelens is set in config."],
         )
 

cvs/parsers/tracelens.py

@@ -64,17 +64,33 @@ def parse(self, run_result: RunResult) -> ParseResult[AortaTraceMetrics]:
         Returns:
             ParseResult containing validated AortaTraceMetrics
         """
+        run_warnings = []
         if not run_result.succeeded:
-            return ParseResult(status=ParseStatus.FAILED, errors=[f"Run did not succeed: {run_result.error_message}"])
+            # A failed/timed-out node no longer discards traces collected from
+            # surviving nodes (see AortaRunner.run()), so a partial run can
+            # still have real data to parse. Only bail out below if there is
+            # nothing on disk to parse.
+            run_warnings.append(f"Run did not succeed: {run_result.error_message}")
+            log.warning(run_warnings[-1])
 
         trace_dir = run_result.get_artifact("torch_traces")
         if not trace_dir:
-            return ParseResult(status=ParseStatus.FAILED, errors=["No torch_traces artifact found in run result"])
+            return ParseResult(
+                status=ParseStatus.FAILED,
+                warnings=run_warnings,
+                errors=["No torch_traces artifact found in run result"],
+            )
 
         if not trace_dir.exists():
-            return ParseResult(status=ParseStatus.FAILED, errors=[f"Trace directory does not exist: {trace_dir}"])
+            return ParseResult(
+                status=ParseStatus.FAILED,
+                warnings=run_warnings,
+                errors=[f"Trace directory does not exist: {trace_dir}"],
+            )
 
-        return self.parse_trace_directory(trace_dir)
+        result = self.parse_trace_directory(trace_dir)
+        result.warnings = run_warnings + list(result.warnings)
+        return result
 
     def parse_trace_directory(self, trace_dir: Path) -> ParseResult[AortaTraceMetrics]:
         """

cvs/parsers/unittests/test_aorta_report.py

@@ -0,0 +1,72 @@
+"""
+Unit tests for cvs.parsers.aorta_report.
+
+Copyright 2025 Advanced Micro Devices, Inc.
+All rights reserved.
+"""
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import pandas as pd
+
+from cvs.parsers.aorta_report import AortaReportParser
+from cvs.parsers.schemas import ParseStatus
+from cvs.runners._base_runner import RunResult, RunStatus
+
+
+def _write_report(individual_dir: Path, rank: int) -> None:
+    individual_dir.mkdir(parents=True, exist_ok=True)
+    df = pd.DataFrame(
+        {
+            "type": ["total_time", "computation_time", "exposed_comm_time", "total_comm_time"],
+            "time ms": [10.0, 6.0, 2.0, 3.0],
+        }
+    )
+    df.to_excel(individual_dir / f"perf_rank{rank}.xlsx", sheet_name="gpu_timeline", index=False)
+
+
+class TestParseGatingOnRunStatus(unittest.TestCase):
+    def setUp(self):
+        self.parser = AortaReportParser()
+
+    def test_failed_run_with_reports_on_disk_still_parses(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            analysis_dir = Path(tmp) / "tracelens_analysis"
+            _write_report(analysis_dir / "individual_reports", 0)
+            run_result = RunResult(
+                status=RunStatus.FAILED,
+                start_time=0,
+                end_time=1,
+                error_message="node b timed out",
+                artifacts={"tracelens_analysis": analysis_dir},
+            )
+            result = self.parser.parse(run_result)
+            self.assertEqual(result.status, ParseStatus.SUCCESS)
+            self.assertIn("Run did not succeed: node b timed out", result.warnings)
+
+    def test_failed_run_with_no_artifact_still_fails(self):
+        run_result = RunResult(status=RunStatus.FAILED, start_time=0, end_time=1, error_message="all nodes died")
+        result = self.parser.parse(run_result)
+        self.assertEqual(result.status, ParseStatus.FAILED)
+        self.assertIn("Run did not succeed: all nodes died", result.warnings)
+        self.assertIn(
+            "No tracelens_analysis artifact found. Ensure analysis.enable_tracelens is set in config.",
+            result.errors,
+        )
+
+    def test_completed_run_parses_without_run_warning(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            analysis_dir = Path(tmp) / "tracelens_analysis"
+            _write_report(analysis_dir / "individual_reports", 0)
+            run_result = RunResult(
+                status=RunStatus.COMPLETED, start_time=0, end_time=1, artifacts={"tracelens_analysis": analysis_dir}
+            )
+            result = self.parser.parse(run_result)
+            self.assertEqual(result.status, ParseStatus.SUCCESS)
+            self.assertEqual(result.warnings, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

cvs/parsers/unittests/test_tracelens.py

@@ -0,0 +1,68 @@
+"""
+Unit tests for cvs.parsers.tracelens.
+
+Copyright 2025 Advanced Micro Devices, Inc.
+All rights reserved.
+"""
+
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+from cvs.parsers.schemas import ParseStatus
+from cvs.parsers.tracelens import TraceLensParser
+from cvs.runners._base_runner import RunResult, RunStatus
+
+
+def _write_trace(trace_dir: Path, rank: int) -> None:
+    rank_dir = trace_dir / f"rank{rank}"
+    rank_dir.mkdir(parents=True, exist_ok=True)
+    trace = {
+        "traceEvents": [
+            {"name": "aten::matmul", "cat": "kernel", "dur": 100},
+        ]
+    }
+    (rank_dir / "trace.json").write_text(json.dumps(trace))
+
+
+class TestParseGatingOnRunStatus(unittest.TestCase):
+    def setUp(self):
+        self.parser = TraceLensParser(use_tracelens=False)
+
+    def test_failed_run_with_traces_on_disk_still_parses(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_dir = Path(tmp) / "torch_traces"
+            _write_trace(trace_dir, 0)
+            run_result = RunResult(
+                status=RunStatus.FAILED,
+                start_time=0,
+                end_time=1,
+                error_message="node b timed out",
+                artifacts={"torch_traces": trace_dir},
+            )
+            result = self.parser.parse(run_result)
+            self.assertEqual(result.status, ParseStatus.SUCCESS)
+            self.assertIn("Run did not succeed: node b timed out", result.warnings)
+
+    def test_failed_run_with_no_artifact_still_fails(self):
+        run_result = RunResult(status=RunStatus.FAILED, start_time=0, end_time=1, error_message="all nodes died")
+        result = self.parser.parse(run_result)
+        self.assertEqual(result.status, ParseStatus.FAILED)
+        self.assertIn("Run did not succeed: all nodes died", result.warnings)
+        self.assertIn("No torch_traces artifact found in run result", result.errors)
+
+    def test_completed_run_parses_without_run_warning(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            trace_dir = Path(tmp) / "torch_traces"
+            _write_trace(trace_dir, 0)
+            run_result = RunResult(
+                status=RunStatus.COMPLETED, start_time=0, end_time=1, artifacts={"torch_traces": trace_dir}
+            )
+            result = self.parser.parse(run_result)
+            self.assertEqual(result.status, ParseStatus.SUCCESS)
+            self.assertEqual(result.warnings, [])
+
+
+if __name__ == "__main__":
+    unittest.main()