add trace merge

Author: drisspg

pyproject.toml

@@ -25,7 +25,8 @@ dependencies = [
     "rich",
     "matplotlib",
     "pandas",
-    "seaborn"
+    "seaborn",
+    "typer",
 ]
 
 [project.optional-dependencies]
@@ -49,6 +50,9 @@ llama = [
     "float8_experimental",
 ]
 
+[project.scripts]
+merge-traces = "transformer_nuggets.utils.merge_traces:app"
+
 # ---------- TOOL CONFIGURATIONS ------------
 [tool.hatch.version]
 source = "vcs"

transformer_nuggets/utils/__init__.py

@@ -13,4 +13,5 @@
 )
 from transformer_nuggets.utils.tracing import LoggingMode, NanInfDetect
 from transformer_nuggets.utils.triton import print_sass
+from transformer_nuggets.utils.merge_traces import merge_traces
 # from transformer_nuggets.utils.model_extraction import extract_attention_data

transformer_nuggets/utils/merge_traces.py

@@ -0,0 +1,71 @@
+"""Merge per-rank Chrome/Perfetto traces into a single multi-process trace."""
+
+from __future__ import annotations
+
+import gzip
+import json
+import sys
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+app = typer.Typer(help="Merge per-rank Chrome/Perfetto traces into one file.")
+
+
+def _open_trace(path: str, mode: str):
+    if path.endswith(".gz"):
+        return gzip.open(path, mode + "t", encoding="utf-8")
+    return open(path, mode, encoding="utf-8")
+
+
+def merge_traces(input_paths: list[str], output_path: str) -> None:
+    merged_events: list[dict] = []
+
+    for rank, path in enumerate(input_paths):
+        with _open_trace(path, "r") as f:
+            data = json.load(f)
+
+        events = data.get("traceEvents", data) if isinstance(data, dict) else data
+
+        merged_events.append(
+            {
+                "ph": "M",
+                "name": "process_name",
+                "pid": rank,
+                "tid": 0,
+                "args": {"name": f"Rank {rank}"},
+            }
+        )
+
+        for ev in events:
+            if ev.get("ph") == "M" and ev.get("name") == "process_name":
+                continue
+            ev["pid"] = rank
+            if "id" in ev and ev.get("ph") in ("s", "t", "f"):
+                ev["id"] = ev["id"] + rank * (1 << 32)
+            merged_events.append(ev)
+
+    with _open_trace(output_path, "w") as f:
+        json.dump({"traceEvents": merged_events}, f, indent=0)
+
+
+@app.command()
+def main(
+    traces: Annotated[list[Path], typer.Argument(help="Input trace files, one per rank, in rank order.")],
+    output: Annotated[Path, typer.Option("-o", "--output", help="Output path.")] = Path(
+        "merged_trace.json.gz"
+    ),
+):
+    """Merge per-rank Chrome/Perfetto traces into a single multi-process Perfetto trace."""
+    for p in traces:
+        if not p.exists():
+            typer.echo(f"Error: {p} not found", err=True)
+            raise typer.Exit(1)
+
+    merge_traces([str(p) for p in traces], str(output))
+    typer.echo(f"Merged {len(traces)} traces -> {output}")
+
+
+if __name__ == "__main__":
+    app()