Add view command

Author: fsimonis

preciceprofiling/perfetto.py

@@ -2,28 +2,34 @@
 from socketserver import TCPServer
 from webbrowser import open_new_tab
 import pathlib
+import json
+import urllib.parse
 
 
-def open_in_perfetto(tracefile: pathlib.Path):
+def open_in_perfetto(tracefile: pathlib.Path, startupCommands=[]):
     print(f"Opening {str(tracefile)} on ui.perfetto.dev")
-    content = tracefile.read_text()
+
+    content = tracefile.read_bytes()
+    contentType = (
+        "text/json" if tracefile.suffix == ".json" else "application/x-protobuf"
+    )
 
     PORT = 9001
     ORIGIN = "https://ui.perfetto.dev"
 
     class TraceHandler(BaseHTTPRequestHandler):
         def do_GET(self):
-            if self.path != "/trace.json":
+            if self.path != "/trace/":
                 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-type", contentType)
             self.send_header("Content-length", str(len(content)))
             self.end_headers()
-            self.wfile.write(content.encode())
+            self.wfile.write(content)
 
         def do_POST(self):
             self.send_error(405)
@@ -35,7 +41,13 @@ def log_message(self, format, *args):
 
     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"
+        address = f"{ORIGIN}/#!/?url=http://127.0.0.1:{PORT}/trace/"
+
+        if startupCommands:
+            print(json.dumps(startupCommands))
+            encoded = urllib.parse.quote(json.dumps(startupCommands))
+            address += f"&startupCommands={encoded}"
+
         open_new_tab(address)
 
         httpd.trace_served = False

preciceprofiling/view.py

@@ -0,0 +1,66 @@
+from preciceprofiling.parsers import addInputArgument
+from preciceprofiling.perfetto import open_in_perfetto
+import argparse
+import sys
+import pathlib
+import json
+
+
+def makeViewParser(add_help: bool = True):
+    view_help = "Open the trace in ui.perfetto.dev"
+    view = argparse.ArgumentParser(description=view_help, add_help=add_help)
+    view.add_argument(
+        "pftrace",
+        nargs="?",
+        type=pathlib.Path,
+        default=pathlib.Path("profiling.pftrace"),
+        help="The perfetto trace file to view",
+    )
+    view.add_argument(
+        "-c",
+        "--commands",
+        type=pathlib.Path,
+        help="File to read startup commands from",
+    )
+    view.add_argument(
+        "-i",
+        "--isolate",
+        type=int,
+        nargs="+",
+        help="Isolate ranks in new workspace",
+    )
+    return view
+
+
+def runView(ns):
+    return viewCommand(ns.pftrace, ns.commands, ns.isolate)
+
+
+def viewCommand(tracefile, commandsFile, isolate):
+    cmds = []
+    if isolate:
+        matcher = "|".join(map(str, isolate))
+        WORKSPACE = "preCICE"
+        cmds.append({"id": "dev.perfetto.CreateWorkspace", "args": [WORKSPACE]})
+        cmds.append(
+            {
+                "id": "dev.perfetto.CopyTracksToWorkspaceByRegexWithAncestors",
+                "args": [f".*Rank ({matcher})$", WORKSPACE, "path"],
+            }
+        )
+        cmds.append({"id": "dev.perfetto.SwitchWorkspace", "args": [WORKSPACE]})
+    if commandsFile:
+        assert commandsFile.exists()
+        cmds.extend(json.loads(commandsFile.read_bytes()))
+    open_in_perfetto(tracefile, startupCommands=cmds)
+    return 0
+
+
+def main():
+    parser = makeViewParser()
+    ns = parser.parse_args()
+    return runView(ns)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pyproject.toml

@@ -41,6 +41,7 @@ precice-profiling-trace = "preciceprofiling.trace:main"
 precice-profiling-export = "preciceprofiling.export:main"
 precice-profiling-histogram = "preciceprofiling.histogram:main"
 precice-profiling-pftrace = "preciceprofiling.pftrace:main"
+precice-profiling-view = "preciceprofiling.view:main"
 
 [tool.setuptools]
 packages=["preciceprofiling"]