From 69fae39c56acef81ccacee037ba97c8288cb2459 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Fri, 14 Aug 2026 15:28:06 -0400 Subject: [PATCH 1/6] feat(aorta): collect per-node traces, survive partial node failure Two problems with profiler artifacts on a multi-node run. First, each node wrote its torch_profiler/ tree to its own filesystem, so the host parser only ever saw the head node's ranks. _collect_multi_node_traces() consolidates every node's trees into /combined_traces/node_/: a local copy when the orchestrator shares the head's filesystem, rsync over SSH otherwise, scp -r where rsync is absent. Per-node failures are logged and skipped rather than aborting the collection. The trace-discovery scan skips anything under combined_traces/ so the stale per-node originals cannot shadow the consolidated set, and seeds trace_mtime from the collected tree - without that seed the existing freshest-trace comparison raises UnboundLocalError once trace_dir can be pre-set. Second, run() returned as soon as any node failed, before collecting anything. One flaky node in a multi-hour run therefore discarded every surviving node's traces, forcing a full rerun or a manual TraceLensParser salvage. Collection and artifact discovery now run unconditionally; the run is still reported FAILED or TIMEOUT with the offending nodes named in error_message, but artifacts come from whatever the survivors produced. Co-Authored-By: Claude --- cvs/runners/aorta.py | 270 +++++++++++++++--- cvs/runners/unittests/test_aorta_multinode.py | 102 +++++++ 2 files changed, 334 insertions(+), 38 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index 691c7230a..f3958deb3 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.""" @@ -284,6 +297,156 @@ 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]) -> 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 -r``. Failures on + individual nodes are logged but do not abort the overall collection; + the returned directory is the best-effort union. + + Returns ``None`` only when nothing could be collected at all. + """ + head = self.head_node + combined_root = self.config.aorta_path / "combined_traces" + 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}" + dest.mkdir(parents=True, exist_ok=True) + + try: + # 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) + # 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) + 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 _copy_local_torch_profilers(self, src_root: Path, dest: Path) -> bool: + """ + Copy any ``torch_profiler/`` trees under ``src_root`` into ``dest``, + preserving the relative path. Used for the head node. + """ + 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 + rel = tp.relative_to(src_root) + target = dest / rel + target.parent.mkdir(parents=True, exist_ok=True) + try: + if target.exists(): + shutil.rmtree(target) + shutil.copytree(tp, target, symlinks=True, dirs_exist_ok=False) + copied = True + except OSError as e: + log.warning(f"Local copy {tp} -> {target} failed: {e}") + return copied + + def _copy_remote_torch_profilers(self, node: str, dest: Path) -> bool: + """ + Pull every ``torch_profiler/`` tree under the remote ``aorta_path`` to + ``dest`` using rsync over SSH. Falls back to ``scp -r`` if rsync is + unavailable. + """ + 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) + + list_cmd = [ + "ssh", + *ssh_opts, + f"{ssh_user}@{node}", + f"find {shlex.quote(remote_root)} -type d -name torch_profiler -not -path '*/combined_traces/*'", + ] + 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 = [p.strip() for p in r.stdout.splitlines() if p.strip()] + if not remote_paths: + return False + + copied = False + rsync_available = ( + subprocess.run( + ["bash", "-lc", "command -v rsync >/dev/null"], + capture_output=True, + ).returncode + == 0 + ) + for rp in remote_paths: + try: + rel = Path(rp).relative_to(remote_root) + except ValueError: + rel = Path(Path(rp).name) + target_parent = dest / rel.parent + target_parent.mkdir(parents=True, exist_ok=True) + + if rsync_available: + cmd = [ + "rsync", + "-az", + "-e", + ssh_cmd, + f"{ssh_user}@{node}:{rp}/", + str(target_parent / rel.name) + "/", + ] + else: + cmd = ["scp", "-r", *ssh_opts, f"{ssh_user}@{node}:{rp}", str(target_parent)] + + log.info(f"[{node}] copying {rp} -> {target_parent / rel.name}") + try: + rr = subprocess.run(cmd, capture_output=True, text=True, timeout=1800) + if rr.returncode == 0: + copied = True + else: + log.warning(f"copy of {rp} from {node} failed (exit {rr.returncode}): {rr.stderr.strip()}") + except (subprocess.TimeoutExpired, FileNotFoundError) as e: + log.warning(f"copy of {rp} 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 +936,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,15 +1065,21 @@ 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) @@ -953,7 +1122,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 +1189,22 @@ 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 mn.collect_traces: + combined = self._collect_multi_node_traces(nodes) + 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,34 +1212,53 @@ def run(self, **kwargs) -> RunResult: nch = self.config.environment.NCCL_MAX_NCHANNELS compute_ch = 256 - nch - trace_dir = None - output_dir = None + trace_dir: Optional[Path] = None + output_dir: Optional[Path] = None + trace_mtime: float = -1.0 - # Search for torch_profiler directories in aorta_path (handles nested dirs like artifacts/*/torch_profiler) + if "torch_traces" in artifacts: + trace_dir = artifacts["torch_traces"] + output_dir = trace_dir.parent + # Multi-node combined_traces should win unless a fresher single-node tree + # is discovered below; seed mtime from this tree so the comparison is valid. + try: + latest_file = max( + trace_dir.glob("**/*"), + key=lambda p: p.stat().st_mtime if p.is_file() else 0, + default=None, + ) + if latest_file is not None and latest_file.is_file(): + trace_mtime = latest_file.stat().st_mtime + else: + trace_mtime = trace_dir.stat().st_mtime + except (ValueError, OSError): + trace_mtime = trace_dir.stat().st_mtime + + # Search for torch_profiler directories in aorta_path (handles nested dirs like artifacts/*/torch_profiler). + # Skip anything inside the combined_traces tree we just collected so the + # original (older) per-node copies don't shadow the consolidated set. + combined_root = self.config.aorta_path / "combined_traces" 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) - 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 - ) - candidate_mtime = ( - latest_file.stat().st_mtime - if latest_file and latest_file.is_file() - else candidate.stat().st_mtime - ) - 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: - trace_dir = candidate - output_dir = candidate.parent - trace_mtime = candidate_mtime + if not candidate.is_dir(): + continue + if combined_traces_in(candidate, combined_root): + continue + try: + latest_file = max( + candidate.glob("**/*"), key=lambda p: p.stat().st_mtime if p.is_file() else 0, default=None + ) + candidate_mtime = ( + latest_file.stat().st_mtime + if latest_file and latest_file.is_file() + else candidate.stat().st_mtime + ) + except (ValueError, OSError): + candidate_mtime = candidate.stat().st_mtime + + if trace_dir is None or candidate_mtime > trace_mtime: + trace_dir = candidate + output_dir = candidate.parent + trace_mtime = candidate_mtime # Required artifact for host-side parsing: torch_traces (parse runs on host, not in container) if trace_dir and trace_dir.exists(): @@ -1123,13 +1316,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, diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index f85371038..192f6bb05 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -10,6 +10,7 @@ All rights reserved. """ +import socket import subprocess import tempfile import threading @@ -27,6 +28,7 @@ AortaMultiNodeConfig, AortaRunner, RcclConfig, + combined_traces_in, ) from cvs.runners.unittests.test_aorta import _make_runner @@ -340,6 +342,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 +455,104 @@ def test_teardown_snapshots_containers_so_concurrent_registration_does_not_crash self.assertTrue(r._teardown_started) +class TestRunPartialNodeFailureStillCollectsTraces(unittest.TestCase): + 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_with(["10.0.0.1", "10.0.0.2"]) + self.assertEqual(result.status, RunStatus.FAILED) + self.assertIn("10.0.0.2", result.error_message) + self.assertEqual(result.get_artifact("torch_traces"), combined_root) + + +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_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 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() + ) + + if __name__ == "__main__": unittest.main() From da3e2f7385c1199958cfb2ed31270824f4ed0059 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Thu, 10 Sep 2026 15:46:26 -0400 Subject: [PATCH 2/6] fix(aorta): stop stale combined_traces copies from shadowing fresh trace dirs combined_traces_in() must be checked against aorta_path directly, not an intermediate combined_root variable that pointed at the wrong directory. --- cvs/runners/aorta.py | 3 +-- cvs/runners/unittests/test_aorta_multinode.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index f3958deb3..72c5552a4 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -1237,11 +1237,10 @@ def run(self, **kwargs) -> RunResult: # Search for torch_profiler directories in aorta_path (handles nested dirs like artifacts/*/torch_profiler). # Skip anything inside the combined_traces tree we just collected so the # original (older) per-node copies don't shadow the consolidated set. - combined_root = self.config.aorta_path / "combined_traces" for candidate in self.config.aorta_path.glob("**/torch_profiler"): if not candidate.is_dir(): continue - if combined_traces_in(candidate, combined_root): + if combined_traces_in(candidate, self.config.aorta_path): continue try: latest_file = max( diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index 192f6bb05..464d4579d 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -480,6 +480,31 @@ def fake_run_single_node(*, node, node_rank, launch_cmd, env): 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): From dd7dcce546f664e49dadd19e92bf0e5899bf04c0 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Thu, 10 Sep 2026 15:46:29 -0400 Subject: [PATCH 3/6] fix(parsers): parse partial-failure runs instead of gating on run_result.succeeded TraceLensParser and AortaReportParser both refused to parse whenever a run was not COMPLETED, undermining trace collection surviving partial node failure. Now they only bail out if there is genuinely nothing on disk to parse, and record the run failure as a warning otherwise. --- cvs/parsers/aorta_report.py | 17 ++++- cvs/parsers/tracelens.py | 24 ++++++-- cvs/parsers/unittests/test_aorta_report.py | 72 ++++++++++++++++++++++ cvs/parsers/unittests/test_tracelens.py | 68 ++++++++++++++++++++ 4 files changed, 174 insertions(+), 7 deletions(-) create mode 100644 cvs/parsers/unittests/test_aorta_report.py create mode 100644 cvs/parsers/unittests/test_tracelens.py 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() From 7926ed06b4fd98c07d5096daac78d4de093f4a70 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Thu, 10 Sep 2026 17:54:00 -0400 Subject: [PATCH 4/6] fix(aorta): trace hygiene, artifact preservation, and report status on partial failure - Clear a node's combined_traces dest before repopulating it, so a node that fails to produce fresh traces this run can't have stale prior-run data mistaken for current-run data. - Resolve the TraceLens/GEMM analysis output_dir to the head node's own torch_profiler tree when trace_dir is the aggregated combined_traces root, instead of passing the whole aorta_path mount. - Preserve already-collected artifacts on RunResult when run() hits an exception after trace collection, instead of dropping them. - Isolate the TraceLens dependency probe in its own try/except so a failure probing for it can't be mistaken for a failure of the run. - Report the run's actual status in the generated benchmark report instead of hardcoding "completed". Co-Authored-By: Claude Sonnet 5 --- cvs/runners/aorta.py | 78 ++++++++++---- cvs/runners/unittests/test_aorta_multinode.py | 102 ++++++++++++++++++ cvs/tests/benchmark/test_aorta.py | 2 +- 3 files changed, 163 insertions(+), 19 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index 72c5552a4..437678195 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -314,6 +314,8 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: Returns ``None`` only when nothing could be collected at all. """ + import shutil + head = self.head_node combined_root = self.config.aorta_path / "combined_traces" try: @@ -325,6 +327,14 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: any_collected = False for rank, node in enumerate(nodes): dest = combined_root / f"node_{rank}" + # Clear any leftover content from a previous run before repopulating, so + # a node that fails to produce traces this run (or is unreachable) can't + # have stale prior-run data silently mixed into this run's metrics. + if dest.exists(): + try: + shutil.rmtree(dest) + except OSError as e: + log.warning(f"Could not clear stale {dest}: {e}") dest.mkdir(parents=True, exist_ok=True) try: @@ -348,6 +358,32 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: 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) -> bool: """ Copy any ``torch_profiler/`` trees under ``src_root`` into ``dest``, @@ -1280,15 +1316,17 @@ def run(self, **kwargs) -> RunResult: # 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}") @@ -1296,13 +1334,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}") @@ -1341,6 +1374,7 @@ def run(self, **kwargs) -> RunResult: stdout=stdout_dict, stderr=stderr_dict, exit_codes=exit_codes, + artifacts=artifacts, error_message=str(e), ) @@ -1368,10 +1402,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 464d4579d..6d54bc166 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -10,6 +10,7 @@ All rights reserved. """ +import shutil import socket import subprocess import tempfile @@ -553,6 +554,85 @@ def test_returns_false_when_no_traces(self): self.assertFalse(runner._copy_local_torch_profilers(root, dest)) +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 @@ -578,6 +658,28 @@ def test_layout_matches_combined_traces_node_rank(self): ).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") + second = runner._collect_multi_node_traces([socket.gethostname()]) + + self.assertIsNone(second) + self.assertFalse(stale_file.exists()) + if __name__ == "__main__": unittest.main() diff --git a/cvs/tests/benchmark/test_aorta.py b/cvs/tests/benchmark/test_aorta.py index b3abdfd9c..85ce2b99b 100644 --- a/cvs/tests/benchmark/test_aorta.py +++ b/cvs/tests/benchmark/test_aorta.py @@ -413,7 +413,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, From e8c0aeab0c29381cd129798234316ecd068d7696 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Thu, 10 Sep 2026 18:36:41 -0400 Subject: [PATCH 5/6] fix(aorta): make trace collection correct under reused output dirs and cluster resize - _collect_multi_node_traces now wipes and recreates combined_traces from scratch each run, so a shrunk cluster no longer leaves a previous larger run's higher-numbered node_ directories for the parser to misread. - _copy_local_torch_profilers/_copy_remote_torch_profilers accept a min_mtime floor and skip torch_profiler trees with no file modified at or after it, so a node that fails before writing new traces no longer has its stale previous-run output copied in as if it were current. run() derives the floor from its own start_time with a clock-skew tolerance. - test_parse_results only prefers container-generated Excel reports on single-node runs, since TraceLens analysis only ever covers the head node's container; on multi-node runs the complete raw-trace parse stays authoritative so metrics aren't silently scoped to the head node. Co-Authored-By: Claude Sonnet 5 --- cvs/runners/aorta.py | 80 ++++++++++--- cvs/runners/unittests/test_aorta_multinode.py | 107 +++++++++++++++++- cvs/tests/benchmark/test_aorta.py | 9 +- 3 files changed, 179 insertions(+), 17 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index 437678195..e6e3b8cd9 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -197,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. @@ -297,7 +302,7 @@ 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]) -> Optional[Path]: + 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. @@ -312,12 +317,30 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: 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 any + torch_profiler tree with no file modified at/after it: a node that is + unreachable or fails before writing new output this run must not have + an older run's torch_profiler tree copied in as if it were current -- + including when the training config reuses the same output_dir across + runs, so the tree itself can't be told apart from a fresh one by path + alone. + 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}") try: combined_root.mkdir(parents=True, exist_ok=True) except OSError as e: @@ -327,14 +350,6 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: any_collected = False for rank, node in enumerate(nodes): dest = combined_root / f"node_{rank}" - # Clear any leftover content from a previous run before repopulating, so - # a node that fails to produce traces this run (or is unreachable) can't - # have stale prior-run data silently mixed into this run's metrics. - if dest.exists(): - try: - shutil.rmtree(dest) - except OSError as e: - log.warning(f"Could not clear stale {dest}: {e}") dest.mkdir(parents=True, exist_ok=True) try: @@ -342,12 +357,12 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: # 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) + 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) + 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}") @@ -358,6 +373,17 @@ def _collect_multi_node_traces(self, nodes: List[str]) -> Optional[Path]: return combined_root if any_collected else None + @staticmethod + def _has_file_newer_than(path: Path, min_mtime: float) -> bool: + """Return ``True`` if ``path`` contains a file modified at/after ``min_mtime``.""" + for f in path.rglob("*"): + try: + if f.is_file() and f.stat().st_mtime >= min_mtime: + return True + except OSError: + continue + return False + 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, @@ -384,10 +410,13 @@ def _resolve_analysis_output_dir(self, trace_dir: Optional[Path], output_dir: Op return None return head_trace.parent - def _copy_local_torch_profilers(self, src_root: Path, dest: Path) -> bool: + 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. + + Trees with no file modified at/after ``min_mtime`` are skipped as + belonging to a previous run (see ``_collect_multi_node_traces``). """ import shutil @@ -397,6 +426,9 @@ def _copy_local_torch_profilers(self, src_root: Path, dest: Path) -> bool: continue if combined_traces_in(tp, src_root): continue + if min_mtime is not None and not self._has_file_newer_than(tp, min_mtime): + log.warning(f"Skipping stale torch_profiler tree from a previous run: {tp}") + continue rel = tp.relative_to(src_root) target = dest / rel target.parent.mkdir(parents=True, exist_ok=True) @@ -409,11 +441,14 @@ def _copy_local_torch_profilers(self, src_root: Path, dest: Path) -> bool: log.warning(f"Local copy {tp} -> {target} failed: {e}") return copied - def _copy_remote_torch_profilers(self, node: str, dest: Path) -> bool: + def _copy_remote_torch_profilers(self, node: str, dest: Path, min_mtime: Optional[float] = None) -> bool: """ Pull every ``torch_profiler/`` tree under the remote ``aorta_path`` to ``dest`` using rsync over SSH. Falls back to ``scp -r`` if rsync is unavailable. + + Trees with no file modified at/after ``min_mtime`` are skipped as + belonging to a previous run (see ``_collect_multi_node_traces``). """ ssh_user = self.config.username remote_root = str(self.config.aorta_path) @@ -423,11 +458,25 @@ def _copy_remote_torch_profilers(self, node: str, dest: Path) -> bool: ssh_opts.extend(["-i", self.config.pkey]) ssh_cmd = "ssh " + " ".join(shlex.quote(p) for p in ssh_opts) + find_dirs = f"find {shlex.quote(remote_root)} -type d -name torch_profiler -not -path '*/combined_traces/*'" + if min_mtime is not None: + # For each candidate dir, only keep it if it has a file newer than + # min_mtime -- a plain directory mtime check is not reliable here + # since training may reuse the same torch_profiler/rank_N/ layout + # across runs, only overwriting file contents in place. + remote_cmd = ( + f"for d in $({find_dirs}); do " + f'if [ -n "$(find "$d" -type f -newermt @{int(min_mtime)} -print -quit)" ]; then echo "$d"; fi; ' + "done" + ) + else: + remote_cmd = find_dirs + list_cmd = [ "ssh", *ssh_opts, f"{ssh_user}@{node}", - f"find {shlex.quote(remote_root)} -type d -name torch_profiler -not -path '*/combined_traces/*'", + remote_cmd, ] try: r = subprocess.run(list_cmd, capture_output=True, text=True, timeout=120) @@ -1237,7 +1286,8 @@ def run(self, **kwargs) -> RunResult: log.error(f"Disaggregated run failed on {len(failed)}/{nnodes} nodes: {failed}") if mn.collect_traces: - combined = self._collect_multi_node_traces(nodes) + 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}") diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index 6d54bc166..9741997f8 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -10,6 +10,7 @@ All rights reserved. """ +import os import shutil import socket import subprocess @@ -476,7 +477,9 @@ def fake_run_single_node(*, node, node_rank, launch_cmd, env): ): result = r.run() - mock_collect.assert_called_once_with(["10.0.0.1", "10.0.0.2"]) + 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) @@ -680,6 +683,108 @@ def test_stale_trace_from_previous_run_is_cleared_before_recollection(self): 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. + first = runner._collect_multi_node_traces([socket.gethostname(), "10.0.0.2", "10.0.0.3"]) + self.assertIsNotNone(first) + self.assertTrue((root / "combined_traces" / "node_2").exists()) + + # 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. + 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 85ce2b99b..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( From 8cd5bdeb51af4108f95c610465bfd005cc377fa7 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Thu, 10 Sep 2026 19:01:11 -0400 Subject: [PATCH 6/6] fix(aorta): filter trace freshness per-file and stop legacy fallback from reintroducing stale traces - _copy_local_torch_profilers/_copy_remote_torch_profilers now filter individual files by mtime instead of skipping or keeping a whole torch_profiler tree, so a reused output dir with a mix of old and new rank files only copies the new ones. - _copy_remote_torch_profilers lists remote files with mtimes via find, filters by freshness, then transfers the selected files in one bulk rsync (falling back to per-file scp), instead of rsyncing/scp-ing whole directories. - run() no longer falls back to scanning aorta_path for any torch_profiler directory (including the hardcoded legacy path name) once multi-node trace collection has already run; that fallback ran unconditionally before and could reintroduce a stale trace that collection had just filtered out. Co-Authored-By: Claude Sonnet 5 --- cvs/runners/aorta.py | 264 ++++++++---------- cvs/runners/unittests/test_aorta_multinode.py | 150 +++++++++- 2 files changed, 269 insertions(+), 145 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index e6e3b8cd9..26d07d641 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -313,17 +313,13 @@ def _collect_multi_node_traces(self, nodes: List[str], min_mtime: Optional[float 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 -r``. Failures on + 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 any - torch_profiler tree with no file modified at/after it: a node that is - unreachable or fails before writing new output this run must not have - an older run's torch_profiler tree copied in as if it were current -- - including when the training config reuses the same output_dir across - runs, so the tree itself can't be told apart from a fresh one by path - alone. + ``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. """ @@ -341,6 +337,7 @@ def _collect_multi_node_traces(self, nodes: List[str], min_mtime: Optional[float 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: @@ -350,9 +347,9 @@ def _collect_multi_node_traces(self, nodes: List[str], min_mtime: Optional[float any_collected = False for rank, node in enumerate(nodes): dest = combined_root / f"node_{rank}" - dest.mkdir(parents=True, exist_ok=True) 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 @@ -373,17 +370,6 @@ def _collect_multi_node_traces(self, nodes: List[str], min_mtime: Optional[float return combined_root if any_collected else None - @staticmethod - def _has_file_newer_than(path: Path, min_mtime: float) -> bool: - """Return ``True`` if ``path`` contains a file modified at/after ``min_mtime``.""" - for f in path.rglob("*"): - try: - if f.is_file() and f.stat().st_mtime >= min_mtime: - return True - except OSError: - continue - return False - 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, @@ -415,8 +401,7 @@ def _copy_local_torch_profilers(self, src_root: Path, dest: Path, min_mtime: Opt Copy any ``torch_profiler/`` trees under ``src_root`` into ``dest``, preserving the relative path. Used for the head node. - Trees with no file modified at/after ``min_mtime`` are skipped as - belonging to a previous run (see ``_collect_multi_node_traces``). + Only files modified at/after ``min_mtime`` are copied when it is set. """ import shutil @@ -426,30 +411,29 @@ def _copy_local_torch_profilers(self, src_root: Path, dest: Path, min_mtime: Opt continue if combined_traces_in(tp, src_root): continue - if min_mtime is not None and not self._has_file_newer_than(tp, min_mtime): - log.warning(f"Skipping stale torch_profiler tree from a previous run: {tp}") - continue - rel = tp.relative_to(src_root) - target = dest / rel - target.parent.mkdir(parents=True, exist_ok=True) - try: - if target.exists(): - shutil.rmtree(target) - shutil.copytree(tp, target, symlinks=True, dirs_exist_ok=False) - copied = True - except OSError as e: - log.warning(f"Local copy {tp} -> {target} failed: {e}") + 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 every ``torch_profiler/`` tree under the remote ``aorta_path`` to - ``dest`` using rsync over SSH. Falls back to ``scp -r`` if rsync is - unavailable. + Pull selected profiler files to ``dest`` using a bulk rsync transfer. + Falls back to individual scp transfers if rsync is unavailable. - Trees with no file modified at/after ``min_mtime`` are skipped as - belonging to a previous run (see ``_collect_multi_node_traces``). + 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) @@ -458,19 +442,15 @@ def _copy_remote_torch_profilers(self, node: str, dest: Path, min_mtime: Optiona ssh_opts.extend(["-i", self.config.pkey]) ssh_cmd = "ssh " + " ".join(shlex.quote(p) for p in ssh_opts) - find_dirs = f"find {shlex.quote(remote_root)} -type d -name torch_profiler -not -path '*/combined_traces/*'" - if min_mtime is not None: - # For each candidate dir, only keep it if it has a file newer than - # min_mtime -- a plain directory mtime check is not reliable here - # since training may reuse the same torch_profiler/rank_N/ layout - # across runs, only overwriting file contents in place. - remote_cmd = ( - f"for d in $({find_dirs}); do " - f'if [ -n "$(find "$d" -type f -newermt @{int(min_mtime)} -print -quit)" ]; then echo "$d"; fi; ' - "done" - ) - else: - remote_cmd = find_dirs + # 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", @@ -488,48 +468,70 @@ def _copy_remote_torch_profilers(self, node: str, dest: Path, min_mtime: Optiona log.warning(f"find on {node} returned {r.returncode}: {r.stderr.strip()}") return False - remote_paths = [p.strip() for p in r.stdout.splitlines() if p.strip()] + 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 - copied = False - rsync_available = ( - subprocess.run( - ["bash", "-lc", "command -v rsync >/dev/null"], - capture_output=True, - ).returncode - == 0 - ) - for rp in remote_paths: + 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: - rel = Path(rp).relative_to(remote_root) - except ValueError: - rel = Path(Path(rp).name) - target_parent = dest / rel.parent - target_parent.mkdir(parents=True, exist_ok=True) - - if rsync_available: - cmd = [ - "rsync", - "-az", - "-e", - ssh_cmd, - f"{ssh_user}@{node}:{rp}/", - str(target_parent / rel.name) + "/", - ] - else: - cmd = ["scp", "-r", *ssh_opts, f"{ssh_user}@{node}:{rp}", str(target_parent)] + 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 - log.info(f"[{node}] copying {rp} -> {target_parent / rel.name}") + copied = False + for relative_path in remote_paths: + target = dest / relative_path try: - rr = subprocess.run(cmd, capture_output=True, text=True, timeout=1800) - if rr.returncode == 0: + 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 {rp} from {node} failed (exit {rr.returncode}): {rr.stderr.strip()}") - except (subprocess.TimeoutExpired, FileNotFoundError) as e: - log.warning(f"copy of {rp} from {node} failed: {e}") - + 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]]: @@ -1167,6 +1169,7 @@ def run(self, **kwargs) -> RunResult: try: 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( @@ -1285,7 +1288,7 @@ def run(self, **kwargs) -> RunResult: partial_failure_message = f"Disaggregated experiment failed on nodes: {sorted(failed.keys())}" log.error(f"Disaggregated run failed on {len(failed)}/{nnodes} nodes: {failed}") - if mn.collect_traces: + 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: @@ -1298,69 +1301,46 @@ def run(self, **kwargs) -> RunResult: nch = self.config.environment.NCCL_MAX_NCHANNELS compute_ch = 256 - nch - trace_dir: Optional[Path] = None - output_dir: Optional[Path] = None + trace_dir = artifacts.get("torch_traces") + output_dir = trace_dir.parent if trace_dir else None trace_mtime: float = -1.0 - if "torch_traces" in artifacts: - trace_dir = artifacts["torch_traces"] - output_dir = trace_dir.parent - # Multi-node combined_traces should win unless a fresher single-node tree - # is discovered below; seed mtime from this tree so the comparison is valid. - try: - latest_file = max( - trace_dir.glob("**/*"), - key=lambda p: p.stat().st_mtime if p.is_file() else 0, - default=None, - ) - if latest_file is not None and latest_file.is_file(): - trace_mtime = latest_file.stat().st_mtime - else: - trace_mtime = trace_dir.stat().st_mtime - except (ValueError, OSError): - trace_mtime = trace_dir.stat().st_mtime - - # Search for torch_profiler directories in aorta_path (handles nested dirs like artifacts/*/torch_profiler). - # Skip anything inside the combined_traces tree we just collected so the - # original (older) per-node copies don't shadow the consolidated set. - for candidate in self.config.aorta_path.glob("**/torch_profiler"): - if not candidate.is_dir(): - continue - if combined_traces_in(candidate, self.config.aorta_path): - continue - try: - latest_file = max( - candidate.glob("**/*"), key=lambda p: p.stat().st_mtime if p.is_file() else 0, default=None - ) - candidate_mtime = ( - latest_file.stat().st_mtime - if latest_file and latest_file.is_file() - else candidate.stat().st_mtime - ) - except (ValueError, OSError): - candidate_mtime = candidate.stat().st_mtime + # 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: + latest_file = max( + candidate.glob("**/*"), key=lambda p: p.stat().st_mtime if p.is_file() else 0, default=None + ) + candidate_mtime = ( + latest_file.stat().st_mtime + if latest_file and latest_file.is_file() + else candidate.stat().st_mtime + ) + except (ValueError, OSError): + candidate_mtime = candidate.stat().st_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 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. diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index 9741997f8..63ce3618b 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -458,6 +458,55 @@ def test_teardown_snapshots_containers_so_concurrent_registration_does_not_crash 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) @@ -525,6 +574,28 @@ def test_returns_false_for_path_outside_root(self): 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) @@ -557,6 +628,75 @@ def test_returns_false_when_no_traces(self): 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 @@ -678,7 +818,8 @@ def test_stale_trace_from_previous_run_is_cleared_before_recollection(self): # 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") - second = runner._collect_multi_node_traces([socket.gethostname()]) + 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()) @@ -691,9 +832,11 @@ def test_shrunk_cluster_does_not_leave_previous_runs_higher_rank_dirs(self): 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. - first = runner._collect_multi_node_traces([socket.gethostname(), "10.0.0.2", "10.0.0.3"]) + 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. @@ -723,7 +866,8 @@ def test_source_tree_older_than_min_mtime_is_skipped(self): 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. - result = runner._collect_multi_node_traces([socket.gethostname()], min_mtime=old_mtime + 1800) + 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())