fix(perfetto): skip full tensor dumps during trace discovery

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:

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
+<rank>/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

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()