Support native .pftrace output in merge_traces

Route merged events through write_track_event_trace when the output path
ends in .pftrace, so merged multi-process traces can be viewed as native
Perfetto TrackEvent protobufs instead of Chrome JSON only.

Author: drisspg

test/test_perfetto.py

@@ -330,3 +330,34 @@ def test_track_event_conversion_keeps_back_to_back_slices_separate():
 
     assert all(not stack for stack in open_stacks.values())
     assert rendered == {"a": (10_000, 20_000), "b": (20_000, 30_000)}
+
+
+def test_merge_traces_writes_native_pftrace(tmp_path):
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import Trace
+
+    from transformer_nuggets.utils.merge_traces import merge_traces
+
+    inputs = []
+    for idx in range(2):
+        path = tmp_path / f"rank{idx}.json"
+        events = [{"ph": "X", "name": f"op{idx}", "pid": 7, "tid": 3, "ts": 100 + idx, "dur": 5}]
+        path.write_text(json.dumps({"traceEvents": events}))
+        inputs.append(str(path))
+
+    output = tmp_path / "merged.pftrace"
+    merge_traces(inputs, str(output), labels=["impl a", "impl b"], align_timestamps=True)
+
+    trace = Trace()
+    trace.ParseFromString(output.read_bytes())
+    process_names = {
+        p.track_descriptor.process.process_name
+        for p in trace.packet
+        if p.HasField("track_descriptor") and p.track_descriptor.HasField("process")
+    }
+    assert {"impl a", "impl b"} <= process_names
+    slice_names = {
+        p.track_event.name
+        for p in trace.packet
+        if p.HasField("track_event") and p.track_event.name
+    }
+    assert {"op0", "op1"} <= slice_names

transformer_nuggets/utils/merge_traces.py

@@ -9,6 +9,8 @@
 
 import typer
 
+from transformer_nuggets.utils.track_event import write_track_event_trace
+
 app = typer.Typer(help="Merge per-rank Chrome/Perfetto traces into one file.")
 
 
@@ -31,6 +33,12 @@ def merge_traces(
     labels: list[str] | None = None,
     align_timestamps: bool = False,
 ) -> None:
+    """Merge Chrome JSON traces into one multi-process trace.
+
+    Output format follows the ``output_path`` suffix: ``.pftrace`` writes a native
+    Perfetto TrackEvent protobuf, anything else writes Chrome JSON (gzipped for
+    ``.gz``).
+    """
     merged_events: list[dict] = []
 
     for idx, path in enumerate(input_paths):
@@ -62,18 +70,26 @@ def merge_traces(
                 ev["id"] = ev["id"] + idx * (1 << 32)
             merged_events.append(ev)
 
-    with _open_trace(output_path, "w") as f:
-        json.dump({"traceEvents": merged_events}, f, indent=0)
+    if output_path.endswith(".pftrace"):
+        write_track_event_trace(output_path, {"traceEvents": merged_events})
+    else:
+        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"
-    ),
+    output: Annotated[
+        Path,
+        typer.Option(
+            "-o",
+            "--output",
+            help="Output path (.pftrace for native Perfetto, .json/.json.gz for Chrome JSON).",
+        ),
+    ] = Path("merged_trace.json.gz"),
     label: Annotated[
         list[str] | None,
         typer.Option("-l", "--label", help="Label for each trace (repeat for each file)."),