diff --git a/cvs/parsers/aorta_report.py b/cvs/parsers/aorta_report.py index 92890aec3..4c3c7784f 100644 --- a/cvs/parsers/aorta_report.py +++ b/cvs/parsers/aorta_report.py @@ -76,13 +76,21 @@ 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") @@ -90,10 +98,13 @@ def parse(self, run_result: RunResult) -> ParseResult[AortaTraceMetrics]: 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."], ) diff --git a/cvs/parsers/tracelens.py b/cvs/parsers/tracelens.py index 36a7e6db3..0da746a08 100644 --- a/cvs/parsers/tracelens.py +++ b/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]: """ diff --git a/cvs/parsers/unittests/test_aorta_report.py b/cvs/parsers/unittests/test_aorta_report.py new file mode 100644 index 000000000..40f178140 --- /dev/null +++ b/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() diff --git a/cvs/parsers/unittests/test_tracelens.py b/cvs/parsers/unittests/test_tracelens.py new file mode 100644 index 000000000..fee2ccc9d --- /dev/null +++ b/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() diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index 691c7230a..26d07d641 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -39,6 +39,19 @@ log = logging.getLogger(__name__) +def combined_traces_in(path: Path, root: Path) -> bool: + """Return True if ``path`` lives under ``root/combined_traces``. + + Used to skip already-collected traces when rescanning the head node so + repeated runs do not nest combined_traces inside itself. + """ + try: + rel = path.relative_to(root) + except ValueError: + return False + return rel.parts and rel.parts[0] == "combined_traces" + + @dataclass class RcclConfig: """RCCL build and runtime configuration.""" @@ -184,6 +197,11 @@ class AortaRunner(BaseRunner): 4. Collect profiling artifacts """ + # Tolerance subtracted from run()'s start time before using it as the trace + # freshness floor in _collect_multi_node_traces, to absorb clock skew between + # the orchestrator and cluster nodes without needing per-node clock sync. + _TRACE_FRESHNESS_SKEW_SECONDS = 60 + def __init__(self, config: AortaConfig): """ Initialize Aorta runner. @@ -284,6 +302,238 @@ def _cleanup_existing_containers(self, client: docker.DockerClient, node: str): except Exception as e: log.warning(f"Error cleaning up container on {node}: {e}") + def _collect_multi_node_traces(self, nodes: List[str], min_mtime: Optional[float] = None) -> Optional[Path]: + """ + Collect torch_profiler trees from every node into a single tree on the + head node and return the parent directory. + + Layout:: + + /combined_traces/node_//torch_profiler/... + + The head node is rsynced locally; non-head nodes are pulled with rsync + over SSH (``rsync -az`` with the configured ``priv_key_file``). When + rsync is unavailable we fall back to ``scp``. Failures on + individual nodes are logged but do not abort the overall collection; + the returned directory is the best-effort union. + + ``min_mtime`` (a ``time.time()``-style epoch), when given, excludes + individual files modified before it. Reusing an output directory + must not mix older ranks or profiler steps into the collected data. + + Returns ``None`` only when nothing could be collected at all. + """ + import shutil + + head = self.head_node + combined_root = self.config.aorta_path / "combined_traces" + # Recreate the whole tree from scratch every run. Clearing only the rank + # directories this run's node list would populate leaves a smaller run's + # combined tree contaminated with higher-numbered node_ directories + # from a previous, larger-cluster run, which the parser would then read + # as if they were current data. + if combined_root.exists(): + try: + shutil.rmtree(combined_root) + except OSError as e: + log.warning(f"Could not clear stale {combined_root}: {e}") + return None + try: + combined_root.mkdir(parents=True, exist_ok=True) + except OSError as e: + log.error(f"Cannot create {combined_root}: {e}") + return None + + any_collected = False + for rank, node in enumerate(nodes): + dest = combined_root / f"node_{rank}" + + try: + dest.mkdir(parents=True, exist_ok=True) + # First pass: copy from the orchestrator's local filesystem. This handles + # the head==orchestrator case and any NFS-shared aorta_path. + found = False + if node == head: + found = self._copy_local_torch_profilers(self.config.aorta_path, dest, min_mtime=min_mtime) + # Pull over SSH for non-head nodes, and also for the head when the + # orchestrator's local fs didn't actually have the head's traces (i.e. + # orchestrator is a separate login node from the head). + if not found: + found = self._copy_remote_torch_profilers(node, dest, min_mtime=min_mtime) + if found: + any_collected = True + log.info(f"Collected traces for node_{rank} ({node}) -> {dest}") + else: + log.warning(f"No torch_profiler artifacts found for node {node} (rank {rank})") + except Exception as e: + log.warning(f"Failed to collect traces for node {node} (rank {rank}): {e}") + + return combined_root if any_collected else None + + def _resolve_analysis_output_dir(self, trace_dir: Optional[Path], output_dir: Optional[Path]) -> Optional[Path]: + """ + Resolve the ``output_dir`` to hand to the container TraceLens/GEMM scripts, + which expect ``output_dir/torch_profiler`` directly. + + When ``trace_dir`` is the aggregated ``combined_traces`` root (multi-node), + ``output_dir`` would otherwise be ``aorta_path`` itself -- the whole mount, + which has no ``torch_profiler`` child of its own -- so resolve to the head + node's own trace tree (``combined_traces/node_0``, rank 0 always being the + head node per ``_collect_multi_node_traces``) instead. Returns ``None`` when + no such tree can be found, so callers skip analysis rather than pass a + directory the scripts cannot use. + """ + if trace_dir is None or output_dir is None: + return None + + combined_root = self.config.aorta_path / "combined_traces" + if trace_dir != combined_root: + return output_dir + + head_trace = next((combined_root / "node_0").glob("**/torch_profiler"), None) + if head_trace is None: + log.warning("No torch_profiler tree found for the head node under combined_traces; skipping analysis") + return None + return head_trace.parent + + def _copy_local_torch_profilers(self, src_root: Path, dest: Path, min_mtime: Optional[float] = None) -> bool: + """ + Copy any ``torch_profiler/`` trees under ``src_root`` into ``dest``, + preserving the relative path. Used for the head node. + + Only files modified at/after ``min_mtime`` are copied when it is set. + """ + import shutil + + copied = False + for tp in src_root.glob("**/torch_profiler"): + if not tp.is_dir(): + continue + if combined_traces_in(tp, src_root): + continue + for source in tp.rglob("*"): + try: + if not source.is_file(): + continue + if min_mtime is not None and source.stat().st_mtime < min_mtime: + continue + target = dest / source.relative_to(src_root) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, target) + copied = True + except OSError as e: + log.warning(f"Local copy of {source} failed: {e}") + return copied + + def _copy_remote_torch_profilers(self, node: str, dest: Path, min_mtime: Optional[float] = None) -> bool: + """ + Pull selected profiler files to ``dest`` using a bulk rsync transfer. + Falls back to individual scp transfers if rsync is unavailable. + + Only files modified at/after ``min_mtime`` are copied when it is set. + """ + import shutil + + ssh_user = self.config.username + remote_root = str(self.config.aorta_path) + + ssh_opts = ["-o", "StrictHostKeyChecking=no", "-o", "BatchMode=yes", "-o", "ConnectTimeout=15"] + if self.config.pkey: + ssh_opts.extend(["-i", self.config.pkey]) + ssh_cmd = "ssh " + " ".join(shlex.quote(p) for p in ssh_opts) + + # Null-delimited records preserve whitespace in paths. Keeping each file's + # mtime in the manifest applies the same cutoff as local collection without + # allowing a fresh file to admit an entire stale profiler tree. + remote_cmd = ( + f"find {shlex.quote(remote_root)} " + f"-path {shlex.quote(str(self.config.aorta_path / 'combined_traces'))} -prune -o " + "-type d -name torch_profiler " + "-exec find {} -type f -printf '%T@ %p\\0' ';' -prune" + ) + + list_cmd = [ + "ssh", + *ssh_opts, + f"{ssh_user}@{node}", + remote_cmd, + ] + try: + r = subprocess.run(list_cmd, capture_output=True, text=True, timeout=120) + except (subprocess.TimeoutExpired, FileNotFoundError) as e: + log.warning(f"Listing remote torch_profiler dirs on {node} failed: {e}") + return False + + if r.returncode != 0: + log.warning(f"find on {node} returned {r.returncode}: {r.stderr.strip()}") + return False + + remote_paths = [] + for record in r.stdout.split("\0"): + if not record: + continue + mtime, separator, source_path = record.partition(" ") + try: + fresh = min_mtime is None or float(mtime) >= min_mtime + except ValueError: + log.warning(f"Invalid trace file timestamp from {node}: {mtime!r}") + continue + try: + path = Path(source_path).relative_to(self.config.aorta_path) + except ValueError: + log.warning(f"Trace file from {node} is outside {remote_root}: {source_path!r}") + continue + if not separator or not path.parts or ".." in path.parts: + log.warning(f"Invalid trace file path from {node}: {source_path!r}") + continue + if fresh: + remote_paths.append(str(path)) + if not remote_paths: + return False + + if shutil.which("rsync"): + cmd = [ + "rsync", + "-az", + "--protect-args", + "--from0", + "--files-from=-", + "-e", + ssh_cmd, + f"{ssh_user}@{node}:{remote_root}/", + str(dest) + "/", + ] + log.info(f"[{node}] copying {len(remote_paths)} profiler files -> {dest}") + try: + result = subprocess.run( + cmd, + input="\0".join(remote_paths) + "\0", + capture_output=True, + text=True, + timeout=1800, + ) + if result.returncode == 0: + return True + log.warning(f"Trace copy from {node} failed (exit {result.returncode}): {result.stderr.strip()}") + except (subprocess.TimeoutExpired, FileNotFoundError) as e: + log.warning(f"Trace copy from {node} failed: {e}") + return False + + copied = False + for relative_path in remote_paths: + target = dest / relative_path + try: + target.parent.mkdir(parents=True, exist_ok=True) + cmd = ["scp", "-p", *ssh_opts, f"{ssh_user}@{node}:{remote_root}/{relative_path}", str(target)] + result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800) + if result.returncode == 0: + copied = True + else: + log.warning(f"Copy of {relative_path} from {node} failed: {result.stderr.strip()}") + except (subprocess.TimeoutExpired, OSError) as e: + log.warning(f"Copy of {relative_path} from {node} failed: {e}") + return copied + def _get_remote_uid_gid(self, node: str) -> Optional[Tuple[int, int]]: """ Get UID and GID for config.username on the given node via SSH. @@ -773,7 +1023,7 @@ def _resolve_master_addr(self) -> str: return self.config.node_vpc_ips.get(self.head_node, self.head_node) def _build_base_env(self) -> Dict[str, str]: - """Build the environment dict exported into the container before launch.""" + """Build the env dict shared by every node's launch.""" env = self.config.environment.to_dict() rccl_path = self.config.rccl.build_path @@ -902,17 +1152,24 @@ def run(self, **kwargs) -> RunResult: ``scripts/multi_node/local_launch.sh`` pattern. ``auto`` picks ``script`` for single-node clusters and ``torchrun`` for multi-node clusters. + + Profiling artifacts (``torch_profiler/`` trees) from every node are + collected into ``/combined_traces/node_/`` on the + head node when running multi-node, and exposed via the + ``torch_traces`` artifact for downstream parsers. """ start_time = time.time() stdout_dict: Dict[str, str] = {} stderr_dict: Dict[str, str] = {} exit_codes: Dict[str, int] = {} artifacts: Dict[str, Path] = {} + partial_failure_status: Optional[RunStatus] = None + partial_failure_message: Optional[str] = None try: - # For now, run on head node only (single node v1) launch_mode = self._resolve_launch_mode() nodes = list(self.config.nodes) + use_collected_traces = launch_mode == "torchrun" and self.config.multi_node.collect_traces if launch_mode == "script" and len(nodes) > 1: return RunResult( @@ -953,7 +1210,7 @@ def run(self, **kwargs) -> RunResult: container, exp_cmd, environment=base_env, - stream=True, # Stream output for real-time feedback + stream=True, ) stdout_dict[node] = output exit_codes[node] = exit_code @@ -1020,17 +1277,23 @@ def run(self, **kwargs) -> RunResult: stdout_dict[n] = out exit_codes[n] = ec + # A failed/timed-out node does not short-circuit trace collection below: + # surviving nodes may hold hours of otherwise-good profiler output, and + # forcing a full rerun (or a manual TraceLensParser salvage) to recover it + # is exactly the failure mode this is meant to avoid. The failure is still + # reported via partial_failure_status/message on the final RunResult. failed = {n: c for n, c in exit_codes.items() if c != 0} if failed: + partial_failure_status = RunStatus.TIMEOUT if not_done else RunStatus.FAILED + partial_failure_message = f"Disaggregated experiment failed on nodes: {sorted(failed.keys())}" log.error(f"Disaggregated run failed on {len(failed)}/{nnodes} nodes: {failed}") - return RunResult( - status=RunStatus.TIMEOUT if not_done else RunStatus.FAILED, - start_time=start_time, - end_time=time.time(), - stdout=stdout_dict, - exit_codes=exit_codes, - error_message=f"Disaggregated experiment failed on nodes: {sorted(failed.keys())}", - ) + + if use_collected_traces: + trace_min_mtime = start_time - self._TRACE_FRESHNESS_SKEW_SECONDS + combined = self._collect_multi_node_traces(nodes, min_mtime=trace_min_mtime) + if combined is not None: + artifacts["torch_traces"] = combined + log.info(f"Combined per-node traces collected at {combined}") # Find torch_profiler directory - Aorta saves traces to output_dir/torch_profiler # The output_dir is configured in the YAML config (e.g., "overlap_debug_repro") @@ -1038,15 +1301,17 @@ def run(self, **kwargs) -> RunResult: nch = self.config.environment.NCCL_MAX_NCHANNELS compute_ch = 256 - nch - trace_dir = None - output_dir = None + trace_dir = artifacts.get("torch_traces") + output_dir = trace_dir.parent if trace_dir else None + trace_mtime: float = -1.0 - # Search for torch_profiler directories in aorta_path (handles nested dirs like artifacts/*/torch_profiler) - for candidate in self.config.aorta_path.glob("**/torch_profiler"): - if candidate.is_dir(): - # Use the most recently modified one (check mtime of rank subdirs or files inside) + # Collection has already filtered individual files. Rescanning sources + # could restore rejected traces or replace the union with just one node. + if not use_collected_traces: + for candidate in self.config.aorta_path.glob("**/torch_profiler"): + if not candidate.is_dir() or combined_traces_in(candidate, self.config.aorta_path): + continue try: - # Get mtime of most recent file in the directory latest_file = max( candidate.glob("**/*"), key=lambda p: p.stat().st_mtime if p.is_file() else 0, default=None ) @@ -1058,45 +1323,40 @@ def run(self, **kwargs) -> RunResult: except (ValueError, OSError): candidate_mtime = candidate.stat().st_mtime - if trace_dir is None: - trace_dir = candidate - output_dir = candidate.parent - trace_mtime = candidate_mtime - elif candidate_mtime > trace_mtime: + if trace_dir is None or candidate_mtime > trace_mtime: trace_dir = candidate output_dir = candidate.parent trace_mtime = candidate_mtime + if trace_dir is None: + output_dir_name = f"nodes1_rccl_develop_commsCh{nch}_computeCh{compute_ch}" + output_dir = self.config.aorta_path / output_dir_name + trace_dir = output_dir / "torch_profiler" + # Required artifact for host-side parsing: torch_traces (parse runs on host, not in container) if trace_dir and trace_dir.exists(): artifacts["torch_traces"] = trace_dir log.info(f"Found trace artifacts at {trace_dir} (host_parse_path will use these)") else: - # Fallback to legacy path format - output_dir_name = f"nodes1_rccl_develop_commsCh{nch}_computeCh{compute_ch}" - output_dir = self.config.aorta_path / output_dir_name - trace_dir = output_dir / "torch_profiler" - if trace_dir.exists(): - artifacts["torch_traces"] = trace_dir - log.info(f"Found trace artifacts at {trace_dir} (host_parse_path will use these)") - else: - log.warning( - "No torch_profiler directory found; host cannot produce benchmark metrics without torch_traces" - ) + trace_dir = None + output_dir = None + log.warning("No trace artifacts collected; host cannot produce benchmark metrics without torch_traces") # Optional container_analysis_path: run TraceLens in container only if enabled and deps present. # Parsing/validation use host venv by default; container reports are consumed when present. # In multi-node mode the head node's container is used; the analysis scripts # operate on traces under aorta_path which (for collected traces) is on the head node. + # + # The TraceLens/GEMM scripts expect output_dir/torch_profiler directly. When + # trace_dir is the aggregated combined_traces root (multi-node), output_dir + # would otherwise be aorta_path itself -- the whole mount, not a directory + # that actually contains torch_profiler -- so resolve it to the head node's + # own trace tree inside combined_traces instead. + analysis_output_dir = self._resolve_analysis_output_dir(trace_dir, output_dir) analysis_container = self._containers.get(self.head_node) - if ( - self.config.analysis.enable_tracelens - and trace_dir - and trace_dir.exists() - and analysis_container is not None - ): + if self.config.analysis.enable_tracelens and analysis_output_dir and analysis_container is not None: log.info("Container TraceLens analysis (optional): attempting in-container report generation") - analysis_result = self._run_tracelens_analysis(analysis_container, output_dir) + analysis_result = self._run_tracelens_analysis(analysis_container, analysis_output_dir) if analysis_result: artifacts["tracelens_analysis"] = analysis_result log.info(f"Container TraceLens analysis completed: {analysis_result}") @@ -1104,13 +1364,8 @@ def run(self, **kwargs) -> RunResult: log.warning("Container TraceLens skipped or failed; host will parse raw traces") # Run GEMM analysis if enabled (optional, same as TraceLens) - if ( - self.config.analysis.enable_gemm_analysis - and trace_dir - and trace_dir.exists() - and analysis_container is not None - ): - gemm_result = self._run_gemm_analysis(analysis_container, output_dir) + if self.config.analysis.enable_gemm_analysis and analysis_output_dir and analysis_container is not None: + gemm_result = self._run_gemm_analysis(analysis_container, analysis_output_dir) if gemm_result: artifacts["gemm_analysis"] = gemm_result log.info(f"GEMM analysis completed: {gemm_result}") @@ -1123,13 +1378,14 @@ def run(self, **kwargs) -> RunResult: break return RunResult( - status=RunStatus.COMPLETED, + status=partial_failure_status or RunStatus.COMPLETED, start_time=start_time, end_time=time.time(), stdout=stdout_dict, stderr=stderr_dict, exit_codes=exit_codes, artifacts=artifacts, + error_message=partial_failure_message, metadata={ "nodes": len(self.config.nodes), "gpus_per_node": self.config.gpus_per_node, @@ -1148,6 +1404,7 @@ def run(self, **kwargs) -> RunResult: stdout=stdout_dict, stderr=stderr_dict, exit_codes=exit_codes, + artifacts=artifacts, error_message=str(e), ) @@ -1175,10 +1432,18 @@ def _run_tracelens_analysis(self, container: Container, output_dir: Path) -> Opt return analysis_dir # Fast dependency check to avoid running long scripts that will fail immediately. - check_cmd = 'python3 -c "import TraceLens"' - check_exit, _ = self._exec_in_container(container, check_cmd) - if check_exit != 0: - log.warning("TraceLens python package not available in container; skipping TraceLens analysis") + # This is best-effort: TraceLens is optional analysis layered on top of + # already-collected artifacts, so a failure probing for it (e.g. the + # container/exec itself misbehaving) must not escape and be mistaken by + # run()'s caller for a failure of the run itself. + try: + check_cmd = 'python3 -c "import TraceLens"' + check_exit, _ = self._exec_in_container(container, check_cmd) + if check_exit != 0: + log.warning("TraceLens python package not available in container; skipping TraceLens analysis") + return None + except Exception as e: + log.warning(f"Could not check for TraceLens in container; skipping TraceLens analysis: {e}") return None # Build the analysis command diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index f85371038..63ce3618b 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -10,6 +10,9 @@ All rights reserved. """ +import os +import shutil +import socket import subprocess import tempfile import threading @@ -27,6 +30,7 @@ AortaMultiNodeConfig, AortaRunner, RcclConfig, + combined_traces_in, ) from cvs.runners.unittests.test_aorta import _make_runner @@ -340,6 +344,7 @@ def fake_run_single_node(*, node, node_rank, launch_cmd, env): with ( patch.object(r, "_run_single_node", side_effect=fake_run_single_node), patch.object(r, "_pick_master_port", return_value=29500), + patch.object(r, "_collect_multi_node_traces", return_value=None), ): start = time.time() result = r.run() @@ -452,5 +457,478 @@ def test_teardown_snapshots_containers_so_concurrent_registration_does_not_crash self.assertTrue(r._teardown_started) +class TestRunPartialNodeFailureStillCollectsTraces(unittest.TestCase): + def test_rejected_old_traces_are_not_restored_by_discovery_or_legacy_fallback(self): + for output_name in ("previous_run", "nodes1_rccl_develop_commsCh112_computeCh144"): + with self.subTest(output_name=output_name), tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + trace = root / output_name / "torch_profiler" / "rank0" / "trace.json" + trace.parent.mkdir(parents=True) + trace.write_text("{}") + os.utime(trace, (1, 1)) + runner = _make_runner(nodes=["head", "worker"], aorta_path=root) + + with ( + patch.object( + runner, + "_run_single_node", + side_effect=lambda *, node, **kw: (node, 1, "failed before profiling"), + ), + patch.object(runner, "_pick_master_port", return_value=29500), + patch.object(runner, "_copy_remote_torch_profilers", return_value=False), + ): + result = runner.run() + + self.assertEqual(result.status, RunStatus.FAILED) + self.assertIsNone(result.get_artifact("torch_traces")) + self.assertEqual(list((root / "combined_traces").rglob("*.json")), []) + + def test_combined_artifact_is_not_replaced_by_a_newer_single_node_source(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + combined = root / "combined_traces" + for rank in (0, 8): + trace = combined / f"node_{rank // 8}" / "torch_profiler" / f"rank{rank}" / "trace.json" + trace.parent.mkdir(parents=True) + trace.write_text("{}") + os.utime(trace, (100, 100)) + source = root / "output" / "torch_profiler" / "rank0" / "trace.json" + source.parent.mkdir(parents=True) + source.write_text("{}") + os.utime(source, (200, 200)) + runner = _make_runner(nodes=["head", "worker"], aorta_path=root) + + with ( + patch.object(runner, "_run_single_node", side_effect=lambda *, node, **kw: (node, 0, "ok")), + patch.object(runner, "_pick_master_port", return_value=29500), + patch.object(runner, "_collect_multi_node_traces", return_value=combined), + ): + result = runner.run() + + self.assertEqual(result.get_artifact("torch_traces"), combined) + + def test_failed_node_does_not_block_trace_collection(self): + with tempfile.TemporaryDirectory() as tmp: + aorta_path = Path(tmp) + combined_root = aorta_path / "combined_traces" + combined_root.mkdir() + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=aorta_path) + + def fake_run_single_node(*, node, node_rank, launch_cmd, env): + if node == "10.0.0.2": + return (node, 1, "boom") + return (node, 0, "ok") + + with ( + patch.object(r, "_run_single_node", side_effect=fake_run_single_node), + patch.object(r, "_pick_master_port", return_value=29500), + patch.object(r, "_collect_multi_node_traces", return_value=combined_root) as mock_collect, + ): + result = r.run() + + mock_collect.assert_called_once() + self.assertEqual(mock_collect.call_args.args[0], ["10.0.0.1", "10.0.0.2"]) + self.assertIn("min_mtime", mock_collect.call_args.kwargs) + self.assertEqual(result.status, RunStatus.FAILED) + self.assertIn("10.0.0.2", result.error_message) + self.assertEqual(result.get_artifact("torch_traces"), combined_root) + + def test_stale_combined_traces_copy_is_not_discovered_as_trace_dir(self): + # combined_traces_in must recognize paths nested under combined_traces + # (e.g. a leftover node_0/torch_profiler/ copy from a prior run in the + # same aorta_path, with this run's own collection disabled/empty) so + # the discovery loop never promotes it as trace_dir -- it's a partial, + # single-node view, not a fresh candidate. + with tempfile.TemporaryDirectory() as tmp: + aorta_path = Path(tmp) + stale_copy = aorta_path / "combined_traces" / "node_0" / "torch_profiler" + stale_copy.mkdir(parents=True) + (stale_copy / "trace.json").write_text("{}") + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=aorta_path) + r.config.multi_node.collect_traces = False + + def fake_run_single_node(*, node, node_rank, launch_cmd, env): + return (node, 0, "ok") + + with ( + patch.object(r, "_run_single_node", side_effect=fake_run_single_node), + patch.object(r, "_pick_master_port", return_value=29500), + ): + result = r.run() + + self.assertIsNone(result.get_artifact("torch_traces")) + + +class TestCombinedTracesIn(unittest.TestCase): + def test_returns_true_when_under_combined_traces(self): + root = Path("/aorta") + self.assertTrue(combined_traces_in(root / "combined_traces" / "node_0" / "torch_profiler", root)) + + def test_returns_false_for_real_run_artifacts(self): + root = Path("/aorta") + self.assertFalse(combined_traces_in(root / "artifacts" / "run1" / "torch_profiler", root)) + + def test_returns_false_for_path_outside_root(self): + root = Path("/aorta") + self.assertFalse(combined_traces_in(Path("/elsewhere/torch_profiler"), root)) + + +class TestCopyLocalTorchProfilers(unittest.TestCase): + def test_only_fresh_files_are_copied_from_a_reused_profiler_tree(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + profiler = root / "output" / "torch_profiler" + for name, mtime in ( + ("rank0/trace_step100.json", 1), + ("rank0/trace_step10.json", 100.5), + ("rank8/trace.json", 1), + ): + trace = profiler / name + trace.parent.mkdir(parents=True, exist_ok=True) + trace.write_text(name) + os.utime(trace, (mtime, mtime)) + dest = root / "combined_traces" / "node_0" + runner = _make_runner(nodes=["head"], aorta_path=root) + + self.assertTrue(runner._copy_local_torch_profilers(root, dest, min_mtime=100.5)) + + copied = [str(path.relative_to(dest)) for path in dest.rglob("*.json")] + self.assertEqual(copied, ["output/torch_profiler/rank0/trace_step10.json"]) + self.assertTrue((profiler / "rank0/trace_step100.json").exists()) + + def test_copies_torch_profiler_trees_and_skips_combined(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + # Real run artifact + (root / "artifacts" / "run1" / "torch_profiler" / "rank_0").mkdir(parents=True) + (root / "artifacts" / "run1" / "torch_profiler" / "rank_0" / "trace.json").write_text("{}") + + # Pre-existing combined traces (must be skipped to avoid recursion) + (root / "combined_traces" / "node_0" / "torch_profiler").mkdir(parents=True) + (root / "combined_traces" / "node_0" / "torch_profiler" / "trace.json").write_text("{}") + + dest = root / "combined_traces" / "node_0_new" + dest.mkdir() + + runner = _make_runner(nodes=["a"], aorta_path=str(root)) + copied = runner._copy_local_torch_profilers(root, dest) + + self.assertTrue(copied) + target = dest / "artifacts" / "run1" / "torch_profiler" / "rank_0" / "trace.json" + self.assertTrue(target.exists(), f"Expected {target} to exist") + # Combined traces tree itself must NOT have been re-copied under dest + self.assertFalse((dest / "combined_traces").exists()) + + def test_returns_false_when_no_traces(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + dest = root / "out" + dest.mkdir() + runner = _make_runner(nodes=["a"], aorta_path=str(root)) + self.assertFalse(runner._copy_local_torch_profilers(root, dest)) + + +class TestCopyRemoteTorchProfilers(unittest.TestCase): + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.root = Path(tmp.name) / "torch_profiler" / "aorta checkout" + self.dest = Path(tmp.name) / "collected" + self.fresh_paths = [ + "run with spaces/torch_profiler/rank0/trace_step10.json", + "run with spaces/torch_profiler/rank1/trace with\nnewline.json", + ] + for name, mtime in ( + *((name, 100.5) for name in self.fresh_paths), + ("run with spaces/torch_profiler/rank0/trace_step100.json", 1), + ("run with spaces/torch_profiler/rank8/trace.json", 1), + ("combined_traces/node_0/torch_profiler/rank0/trace.json", 200), + ("unrelated.json", 200), + ): + path = self.root / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(name) + os.utime(path, (mtime, mtime)) + self.runner = _make_runner(nodes=["head", "worker"], aorta_path=self.root) + self.run_process = subprocess.run + + def _mock_transport(self, cmd, **kwargs): + if cmd[0] == "ssh": + return self.run_process(["sh", "-c", cmd[-1]], capture_output=True, text=True) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + def test_rsync_receives_only_fresh_files_with_whitespace_preserved(self): + with ( + patch.object(aorta_mod.subprocess, "run", side_effect=self._mock_transport) as transport, + patch("shutil.which", return_value="/usr/bin/rsync"), + ): + self.assertTrue(self.runner._copy_remote_torch_profilers("worker", self.dest, min_mtime=100.5)) + + self.assertEqual(transport.call_count, 2) + transfer = transport.call_args + self.assertEqual(set(transfer.kwargs["input"].split("\0")[:-1]), set(self.fresh_paths)) + cmd = transfer.args[0] + self.assertIn("--from0", cmd) + self.assertIn("--files-from=-", cmd) + self.assertIn("--protect-args", cmd) + self.assertEqual(cmd[-2:], [f"testuser@worker:{self.root}/", str(self.dest) + "/"]) + + def test_scp_fallback_copies_only_selected_files(self): + with ( + patch.object(aorta_mod.subprocess, "run", side_effect=self._mock_transport) as transport, + patch("shutil.which", return_value=None), + ): + self.assertTrue(self.runner._copy_remote_torch_profilers("worker", self.dest, min_mtime=100.5)) + + transfers = [call.args[0] for call in transport.call_args_list[1:]] + self.assertTrue(all(cmd[0] == "scp" for cmd in transfers)) + self.assertEqual({cmd[-2] for cmd in transfers}, {f"testuser@worker:{self.root}/{p}" for p in self.fresh_paths}) + self.assertEqual({cmd[-1] for cmd in transfers}, {str(self.dest / p) for p in self.fresh_paths}) + + def test_no_fresh_remote_files_does_not_start_a_transfer(self): + with patch.object(aorta_mod.subprocess, "run", side_effect=self._mock_transport) as transport: + self.assertFalse(self.runner._copy_remote_torch_profilers("worker", self.dest, min_mtime=300)) + transport.assert_called_once() + + def test_failed_listing_is_reported_without_starting_a_transfer(self): + result = subprocess.CompletedProcess([], 255, stdout="", stderr="connection failed") + with patch.object(aorta_mod.subprocess, "run", return_value=result) as transport: + self.assertFalse(self.runner._copy_remote_torch_profilers("worker", self.dest, min_mtime=100)) + transport.assert_called_once() + + +class TestRunTracelensAnalysisDependencyCheck(unittest.TestCase): + def test_dependency_check_exception_is_caught_not_raised(self): + # TraceLens analysis is optional, best-effort post-processing. An + # exception while merely probing whether the package is importable + # (e.g. the container exec plumbing itself failing) must not escape + # and be mistaken by run()'s caller for the run itself having failed. + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + with patch.object(r, "_exec_in_container", side_effect=RuntimeError("docker exec blew up")): + result = r._run_tracelens_analysis(Mock(), Path("/tmp/aorta/some_run")) + self.assertIsNone(result) + + def test_dependency_missing_skips_without_raising(self): + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + with patch.object(r, "_exec_in_container", return_value=(1, "ModuleNotFoundError")): + result = r._run_tracelens_analysis(Mock(), Path("/tmp/aorta/some_run")) + self.assertIsNone(result) + + +class TestRunExceptionPreservesArtifacts(unittest.TestCase): + def test_exception_after_trace_collection_still_returns_collected_artifacts(self): + # A failure in later, optional post-processing (e.g. TraceLens analysis) + # must not discard torch_traces and other artifacts already collected + # earlier in run() -- those are exactly what a partially-failed run + # needs to be salvageable. + with tempfile.TemporaryDirectory() as tmp: + aorta_path = Path(tmp) + combined_root = aorta_path / "combined_traces" + (combined_root / "node_0").mkdir(parents=True) + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=aorta_path) + + def fake_run_single_node(*, node, node_rank, launch_cmd, env): + return (node, 0, "ok") + + with ( + patch.object(r, "_run_single_node", side_effect=fake_run_single_node), + patch.object(r, "_pick_master_port", return_value=29500), + patch.object(r, "_collect_multi_node_traces", return_value=combined_root), + patch.object(r, "_resolve_analysis_output_dir", side_effect=RuntimeError("boom")), + ): + result = r.run() + + self.assertEqual(result.status, RunStatus.FAILED) + self.assertEqual(result.get_artifact("torch_traces"), combined_root) + + +class TestResolveAnalysisOutputDir(unittest.TestCase): + def test_non_combined_trace_dir_is_returned_unchanged(self): + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + trace_dir = Path("/tmp/aorta/some_run/torch_profiler") + output_dir = trace_dir.parent + self.assertEqual(r._resolve_analysis_output_dir(trace_dir, output_dir), output_dir) + + def test_combined_trace_dir_resolves_to_head_nodes_own_tree(self): + # trace_dir/output_dir as computed in run() for the combined multi-node + # case: output_dir is aorta_path itself, which has no torch_profiler + # child of its own and is useless to the analysis scripts. + with tempfile.TemporaryDirectory() as tmp: + aorta_path = Path(tmp) + head_trace = aorta_path / "combined_traces" / "node_0" / "artifacts" / "torch_profiler" + head_trace.mkdir(parents=True) + + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=aorta_path) + trace_dir = aorta_path / "combined_traces" + resolved = r._resolve_analysis_output_dir(trace_dir, aorta_path) + + self.assertEqual(resolved, head_trace.parent) + + def test_combined_trace_dir_with_no_head_tree_returns_none(self): + with tempfile.TemporaryDirectory() as tmp: + aorta_path = Path(tmp) + (aorta_path / "combined_traces" / "node_0").mkdir(parents=True) + + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=aorta_path) + trace_dir = aorta_path / "combined_traces" + resolved = r._resolve_analysis_output_dir(trace_dir, aorta_path) + + self.assertIsNone(resolved) + + +class TestCollectMultiNodeTracesHeadOnly(unittest.TestCase): + """ + End-to-end happy path for trace collection where every node is the head + (no SSH involved) so we can exercise the directory layout logic without a + real cluster. + """ + + def test_layout_matches_combined_traces_node_rank(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "artifacts" / "torch_profiler" / "rank_0").mkdir(parents=True) + (root / "artifacts" / "torch_profiler" / "rank_0" / "trace.json").write_text("{}") + + # Single-node "cluster" so the head-node fast path is used for both ranks. + runner = _make_runner(nodes=[socket.gethostname()], aorta_path=str(root)) + result = runner._collect_multi_node_traces([socket.gethostname()]) + + self.assertIsNotNone(result) + self.assertEqual(result, root / "combined_traces") + self.assertTrue( + ( + root / "combined_traces" / "node_0" / "artifacts" / "torch_profiler" / "rank_0" / "trace.json" + ).exists() + ) + + def test_stale_trace_from_previous_run_is_cleared_before_recollection(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "artifacts" / "torch_profiler" / "rank_0").mkdir(parents=True) + (root / "artifacts" / "torch_profiler" / "rank_0" / "trace.json").write_text("run1") + + runner = _make_runner(nodes=[socket.gethostname()], aorta_path=str(root)) + first = runner._collect_multi_node_traces([socket.gethostname()]) + self.assertIsNotNone(first) + stale_file = root / "combined_traces" / "node_0" / "artifacts" / "torch_profiler" / "rank_0" / "trace.json" + self.assertEqual(stale_file.read_text(), "run1") + + # Second run: this node produces no new torch_profiler output this + # time (e.g. training crashed before profiling started). The prior + # run's copy must not still be sitting in combined_traces, where + # it would be mistaken for this run's data. + shutil.rmtree(root / "artifacts" / "torch_profiler") + with patch.object(runner, "_copy_remote_torch_profilers", return_value=False): + second = runner._collect_multi_node_traces([socket.gethostname()]) + + self.assertIsNone(second) + self.assertFalse(stale_file.exists()) + + def test_shrunk_cluster_does_not_leave_previous_runs_higher_rank_dirs(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "artifacts" / "torch_profiler" / "rank_0").mkdir(parents=True) + (root / "artifacts" / "torch_profiler" / "rank_0" / "trace.json").write_text("{}") + + runner = _make_runner(nodes=[socket.gethostname()], aorta_path=str(root)) + # First run had 3 nodes; node_2 belongs to a node no longer in this run. + with patch.object(runner, "_copy_remote_torch_profilers", return_value=False) as remote_copy: + first = runner._collect_multi_node_traces([socket.gethostname(), "worker1", "worker2"]) + self.assertIsNotNone(first) + self.assertTrue((root / "combined_traces" / "node_2").exists()) + self.assertEqual(remote_copy.call_count, 2) + + # This run only has one node -- the old node_2 must not survive, or a + # parser walking combined_traces would still see its stale rank data. + second = runner._collect_multi_node_traces([socket.gethostname()]) + self.assertIsNotNone(second) + self.assertFalse((root / "combined_traces" / "node_2").exists()) + self.assertFalse((root / "combined_traces" / "node_1").exists()) + self.assertTrue((root / "combined_traces" / "node_0").exists()) + + +class TestCollectMultiNodeTracesSourceFreshness(unittest.TestCase): + """ + ``min_mtime`` lets collection tell a previous run's leftover torch_profiler + output apart from this run's, even when the training config reuses the + same output_dir across runs (so the path alone can't tell them apart). + """ + + def test_source_tree_older_than_min_mtime_is_skipped(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + trace_file = root / "artifacts" / "torch_profiler" / "rank_0" / "trace.json" + trace_file.parent.mkdir(parents=True) + trace_file.write_text("stale") + old_mtime = trace_file.stat().st_mtime - 3600 + os.utime(trace_file, (old_mtime, old_mtime)) + + runner = _make_runner(nodes=[socket.gethostname()], aorta_path=str(root)) + # This run "started" after the file above was written, so it must + # not be mistaken for this run's data. + with patch.object(runner, "_copy_remote_torch_profilers", return_value=False): + result = runner._collect_multi_node_traces([socket.gethostname()], min_mtime=old_mtime + 1800) + + self.assertIsNone(result) + self.assertFalse((root / "combined_traces" / "node_0" / "artifacts").exists()) + + def test_source_tree_newer_than_min_mtime_is_collected(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + trace_file = root / "artifacts" / "torch_profiler" / "rank_0" / "trace.json" + trace_file.parent.mkdir(parents=True) + trace_file.write_text("fresh") + run_start = trace_file.stat().st_mtime - 60 + + runner = _make_runner(nodes=[socket.gethostname()], aorta_path=str(root)) + result = runner._collect_multi_node_traces([socket.gethostname()], min_mtime=run_start) + + self.assertIsNotNone(result) + self.assertTrue( + ( + root / "combined_traces" / "node_0" / "artifacts" / "torch_profiler" / "rank_0" / "trace.json" + ).exists() + ) + + def test_no_min_mtime_collects_regardless_of_age(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + trace_file = root / "artifacts" / "torch_profiler" / "rank_0" / "trace.json" + trace_file.parent.mkdir(parents=True) + trace_file.write_text("old but no floor given") + old_mtime = trace_file.stat().st_mtime - 3600 + os.utime(trace_file, (old_mtime, old_mtime)) + + runner = _make_runner(nodes=[socket.gethostname()], aorta_path=str(root)) + result = runner._collect_multi_node_traces([socket.gethostname()]) + + self.assertIsNotNone(result) + + +class TestRunPassesTraceFreshnessFloor(unittest.TestCase): + def test_run_passes_min_mtime_derived_from_start_time(self): + with tempfile.TemporaryDirectory() as tmp: + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=Path(tmp)) + + def fake_run_single_node(*, node, node_rank, launch_cmd, env): + return (node, 0, "ok") + + before = time.time() + with ( + patch.object(r, "_run_single_node", side_effect=fake_run_single_node), + patch.object(r, "_pick_master_port", return_value=29500), + patch.object(r, "_collect_multi_node_traces", return_value=None) as mock_collect, + ): + r.run() + after = time.time() + + min_mtime = mock_collect.call_args.kwargs["min_mtime"] + # Derived from run()'s own start_time (bracketed by before/after), + # offset by exactly the configured skew tolerance -- not some + # unrelated or hardcoded value. + self.assertGreaterEqual(min_mtime, before - r._TRACE_FRESHNESS_SKEW_SECONDS - 1) + self.assertLessEqual(min_mtime, after - r._TRACE_FRESHNESS_SKEW_SECONDS + 1) + + if __name__ == "__main__": unittest.main() diff --git a/cvs/tests/benchmark/test_aorta.py b/cvs/tests/benchmark/test_aorta.py index b3abdfd9c..cd4f35172 100644 --- a/cvs/tests/benchmark/test_aorta.py +++ b/cvs/tests/benchmark/test_aorta.py @@ -284,12 +284,19 @@ def test_parse_results(self, aorta_runner_config, validated_aorta_config): run_result = TestAortaBenchmark.run_result trace_dir = run_result.get_artifact("torch_traces") + # Container TraceLens analysis only ever runs against the head node's own + # container and trace tree (see AortaRunner._resolve_analysis_output_dir), + # so its Excel reports cover only the head node's ranks. On a multi-node + # run that is incomplete, so the complete raw-trace parse below must stay + # authoritative; only a single-node run's reports cover the whole cluster. + is_multi_node = len(aorta_runner_config.nodes) > 1 + # Optional: try container-generated Excel reports first (if present and valid) analysis_dir = run_result.get_artifact("tracelens_analysis") or ( trace_dir.parent / "tracelens_analysis" if trace_dir else None ) has_valid_reports = False - if analysis_dir and analysis_dir.exists(): + if not is_multi_node and analysis_dir and analysis_dir.exists(): reports_dir = analysis_dir / "individual_reports" if reports_dir.exists(): report_files = list(reports_dir.glob("perf_rank*.xlsx")) or list( @@ -413,7 +420,7 @@ def test_generate_report(self, aorta_runner_config): # Build report report = { - "status": "completed", + "status": run_result.status.value if run_result else "unknown", "duration_seconds": run_result.duration_seconds if run_result else 0, "cluster": { "nodes": result.num_nodes,