Add flag to open trace in perfetto

Author: fsimonis

preciceprofiling/perfetto.py

@@ -0,0 +1,44 @@
+from http.server import BaseHTTPRequestHandler
+from socketserver import TCPServer
+from webbrowser import open_new_tab
+import pathlib
+
+
+def open_in_perfetto(tracefile: pathlib.Path):
+    print(f"Opening {str(tracefile)} on ui.perfetto.dev")
+    content = tracefile.read_text()
+
+    PORT = 9001
+    ORIGIN = "https://ui.perfetto.dev"
+
+    class TraceHandler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            if self.path != "/trace.json":
+                self.send_error(405)
+
+            self.server.trace_served = True
+            self.send_response(200)
+            self.send_header("Access-Control-Allow-Origin", ORIGIN)
+            self.send_header("Cache-Control", "no-cache")
+            self.send_header("Content-type", "text/json")
+            self.send_header("Content-length", str(len(content)))
+            self.end_headers()
+            self.wfile.write(content.encode())
+
+        def do_POST(self):
+            self.send_error(405)
+
+        def log_message(self, format, *args):
+            pass
+
+    # We reuse the HTTP+RPC port because it's the only one allowed by the CSP.
+
+    TCPServer.allow_reuse_address = True
+    with TCPServer(("127.0.0.1", PORT), TraceHandler) as httpd:
+        address = f"{ORIGIN}/#!/?url=http://127.0.0.1:{PORT}/trace.json"
+        open_new_tab(address)
+
+        httpd.trace_served = False
+        httpd.allow_origin = ORIGIN
+        while not httpd.trace_served:
+            httpd.handle_request()

preciceprofiling/trace.py

@@ -2,6 +2,8 @@
 import orjson
 import argparse
 import sys
+import pathlib
+from preciceprofiling.perfetto import open_in_perfetto
 
 
 def makeTraceParser(add_help: bool = True):
@@ -15,7 +17,17 @@ def makeTraceParser(add_help: bool = True):
         help="The profiling file to process",
     )
     trace.add_argument(
-        "-o", "--output", default="trace.json", help="The resulting trace file"
+        "-o",
+        "--output",
+        default="trace.json",
+        type=pathlib.Path,
+        help="The resulting trace file",
+    )
+    trace.add_argument(
+        "-w",
+        "--web",
+        action="store_true",
+        help="Open resulting trace in ui.perfetto.dev",
     )
     trace.add_argument(
         "-l", "--limit", type=int, metavar="n", help="Select the first n ranks"
@@ -27,10 +39,10 @@ def makeTraceParser(add_help: bool = True):
 
 
 def runTrace(ns):
-    return traceCommand(ns.profilingfile, ns.output, ns.rank, ns.limit)
+    return traceCommand(ns.profilingfile, ns.output, ns.rank, ns.limit, ns.web)
 
 
-def traceCommand(profilingfile, outfile, rankfilter, limit):
+def traceCommand(profilingfile, outfile, rankfilter, limit, web):
     run = Run(profilingfile)
     selection = (
         set()
@@ -39,8 +51,10 @@ def traceCommand(profilingfile, outfile, rankfilter, limit):
     )
     traces = run.toTrace(selection)
     print(f"Writing to {outfile}")
-    with open(outfile, "wb") as outfile:
-        outfile.write(orjson.dumps(traces))
+    outfile.write_bytes(orjson.dumps(traces))
+
+    if web:
+        open_in_perfetto(outfile)
     return 0
 
 

tests/test_examples.py

@@ -39,7 +39,7 @@ def run_case(case: pathlib.Path, cwd: pathlib.Path, useDir: bool):
     assert export.exists()
 
     print("--- Trace")
-    assert traceCommand(profiling, trace, unit, None) == 0
+    assert traceCommand(profiling, trace, unit, None, False) == 0
     assert trace.exists()
 
     participants = (