diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index f8061bbfa..6572fe6fd 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.""" @@ -280,6 +293,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. @@ -721,7 +884,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 @@ -833,15 +996,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) @@ -874,8 +1043,6 @@ def run(self, **kwargs) -> RunResult: error_message=f"No container found for {node}", ) - # Pass the base config file to the experiment script - # launch_rocm.sh expects: CONFIG=${1:-default.yaml} config_path = f"{self.config.container_mount_path}/{self.config.base_config}" exp_cmd = f"bash {self.config.container_mount_path}/{self.config.experiment_script} {config_path}" log.info(f"Running experiment: {exp_cmd}") @@ -885,7 +1052,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 @@ -952,17 +1119,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") @@ -970,34 +1142,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(): @@ -1055,13 +1246,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 e267d28ef..f03f19230 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() @@ -411,5 +414,104 @@ def test_container_launched_before_cancel_is_registered_normally(self): fake_container.stop.assert_not_called() +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()