diff --git a/opguard/trace_export/perfetto/timeline.py b/opguard/trace_export/perfetto/timeline.py index e2b6f17..6b20a3e 100644 --- a/opguard/trace_export/perfetto/timeline.py +++ b/opguard/trace_export/perfetto/timeline.py @@ -1,4 +1,5 @@ from __future__ import annotations +import os from pathlib import Path from typing import Dict, List @@ -14,16 +15,26 @@ def sort_records(seq: List[dict]) -> List[dict]: keyed.sort(key=lambda x: (x[0], x[1], x[2])) return [r for *_, r in keyed] +def _trace_files(base: Path): + for root, dirs, files in os.walk(base): + # ops/full contains tensor dumps, not operator event records. + if Path(root).name == "ops": + dirs[:] = [name for name in dirs if name != "full"] + for name in files: + path = Path(root) / name + if path.is_file(): + yield path + + def list_pairs(base1: Path, base2: Path): base1, base2 = base1.resolve(), base2.resolve() pairs: List[FilePair] = []; only1: List[Path] = []; only2: List[Path] = [] b2map: Dict[str, Path] = {} n2 = len(str(base2)) + 1 - for q in base2.rglob("*"): - if q.is_file(): b2map[str(q.resolve())[n2:]] = q.resolve() + for q in _trace_files(base2): + b2map[str(q.resolve())[n2:]] = q.resolve() n1 = len(str(base1)) + 1; matched = set() - for p in base1.rglob("*"): - if not p.is_file(): continue + for p in _trace_files(base1): rel = str(p.resolve())[n1:] q = b2map.get(rel) if q: diff --git a/tests/test_perfetto_full_tensor_export.py b/tests/test_perfetto_full_tensor_export.py new file mode 100644 index 0000000..87a8371 --- /dev/null +++ b/tests/test_perfetto_full_tensor_export.py @@ -0,0 +1,89 @@ +"""Perfetto export must tolerate raw tensors from OpGuard's full dump mode. + +TraceBuffer.write_tensor stores dump_tensors=True payloads as plain tensors under +/ops/full/. These files must be pruned before operator records are loaded, +while ordinary events and metadata keep their existing handling. +""" +import gzip +import json +from pathlib import Path + +import torch + +from opguard.trace_export.perfetto.build import build_perfetto + + +def call(output): + return {"function": "aten::mm", "fname": "/checkout/model.py", "line_num": 12, + "gpu_rank": 0, "counter": 1, + "inputs": {"args": [1], "kwargs": {}}, "outputs": [output]} + + +def rank_dir(root, output): + """Mimic one rank directory produced with dump_tensors=True.""" + events = root / "ops" / "aten__mm" + events.mkdir(parents=True) + (events / "aten__mm_line12.jsonl").write_text(json.dumps(call(output)) + "\n") + full = root / "ops" / "full" / "aten::mm" / "gsNA_mbNA_layerNA" + full.mkdir(parents=True) + torch.save(torch.ones(2, 2), full / "model.py_line12_input_args_0.pt") + torch.save(torch.full((2, 2), float(output)), full / "model.py_line12_output_0.pt") + (root / "metadata.json").write_text('{"source_files": {}}') + return root + + +def export(tmp_path, left_output, right_output): + left = rank_dir(tmp_path / "a", left_output) + right = rank_dir(tmp_path / "b", right_output) + target = tmp_path / "trace.json.gz" + build_perfetto(str(left), str(right), str(target)) + with gzip.open(target, "rt") as stream: + doc = json.load(stream) + return [e for e in doc["traceEvents"] if e.get("cat") == "opguard" and e.get("ph") == "X"] + + +def test_full_dump_tensors_are_not_operator_records(tmp_path): + ops = export(tmp_path, 0, 0) + assert len(ops) == 2 + ids = {e["args"]["record_id"] for e in ops} + for event in ops: + assert event["args"]["correspondence"]["peer_record_id"] in ids + + +def test_full_dump_tensors_do_not_hide_a_real_difference(tmp_path): + ops = export(tmp_path, 0, 1) + assert len(ops) == 2 + assert {e["args"]["status"] for e in ops} == {"diff"} + + +def test_full_dump_tensors_never_call_torch_load(tmp_path, monkeypatch): + def unexpected_load(*args, **kwargs): + raise AssertionError("Full tensor dumps must not reach torch.load") + + monkeypatch.setattr(torch, "load", unexpected_load) + assert len(export(tmp_path, 0, 0)) == 2 + + +def test_legacy_pt_operator_records_are_still_loaded(tmp_path, monkeypatch): + left = rank_dir(tmp_path / "a", 0) + right = rank_dir(tmp_path / "b", 0) + legacy_paths = {root / "ops" / "aten__mm" / "legacy.pt" for root in (left, right)} + for path in legacy_paths: + torch.save(dict(call(0), counter=2), path) + original_load = torch.load + loaded = [] + + def checked_load(path, *args, **kwargs): + path = Path(path) + assert "full" not in path.parts + loaded.append(path) + return original_load(path, *args, **kwargs) + + monkeypatch.setattr(torch, "load", checked_load) + target = tmp_path / "trace.json.gz" + build_perfetto(str(left), str(right), str(target)) + with gzip.open(target, "rt") as stream: + operations = [e for e in json.load(stream)["traceEvents"] + if e.get("cat") == "opguard" and e.get("ph") == "X"] + assert len(operations) == 4 + assert set(loaded) == legacy_paths diff --git a/tests/test_perfetto_skip_full_dir.py b/tests/test_perfetto_skip_full_dir.py new file mode 100644 index 0000000..caceb66 --- /dev/null +++ b/tests/test_perfetto_skip_full_dir.py @@ -0,0 +1,45 @@ +"""Full tensor dumps must be pruned before event files are loaded.""" +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from opguard.trace_export.perfetto.timeline import list_pairs + + +class SkipFullDirectoryTests(unittest.TestCase): + def test_prune_full_dumps_preserving_other_pairing(self): + with tempfile.TemporaryDirectory() as temp: + for prefix in ("", "exec_trace/gpu_000/"): + bases = [Path(temp) / str(len(prefix)) / side for side in ("left", "right")] + common = {prefix + "ops/mm/events.jsonl", prefix + "ops/mm/legacy.pt", + "reports/full/metadata.json"} + for base in bases: + for relative in common | {prefix + "ops/full/mm/shared.pt"}: + path = base / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + for base, side in zip(bases, ("left", "right")): + (base / (prefix + f"ops/full/mm/{side}.pt")).touch() + (base / (prefix + f"ops/mm/{side}.jsonl")).touch() + visited = [] + real_walk = os.walk + + def walk(*args, **kwargs): + for entry in real_walk(*args, **kwargs): + visited.append(Path(entry[0])) + yield entry + + with patch("opguard.trace_export.perfetto.timeline.os.walk", walk): + pairs, left, right = list_pairs(*bases) + self.assertEqual({pair.rel for pair in pairs}, common) + self.assertEqual(left, [(bases[0] / (prefix + "ops/mm/left.jsonl")).resolve()]) + self.assertEqual(right, [(bases[1] / (prefix + "ops/mm/right.jsonl")).resolve()]) + for base in bases: + full = base / (prefix + "ops/full") + self.assertFalse(any(full == p or full in p.parents for p in visited)) + + +if __name__ == "__main__": + unittest.main()