Switch traversal benchmark to pyinstrument/speedscope

inducer · Sep 30, 2024 · 3eb71bb · 3eb71bb
1 parent b455b43
commit 3eb71bb
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/experiments/traversal-benchmark.py b/experiments/traversal-benchmark.py
@@ -128,9 +128,11 @@ def main():
         t_end = time()
         print(f"Took: {t_end-t_start} secs.")
     else:
-        import vmprof
-        with open("test.prof", "w+b") as fd:
-            vmprof.enable(fd.fileno())
+        import pyinstrument
+        from pyinstrument.renderers import SpeedscopeRenderer
+        prof = pyinstrument.Profiler()
+        with prof:
             for _ in range(10_000):
                 main()
-            vmprof.disable()
+        with open("ss.json", "w") as outf:
+            outf.write(prof.output(SpeedscopeRenderer(show_all=True)))