From d6796cec113f65dc53872480d7ce7dcf1ba40771 Mon Sep 17 00:00:00 2001
From: "Maarten A. Breddels" <maartenbreddels@gmail.com>
Date: Tue, 15 Dec 2020 12:06:36 +0100
Subject: [PATCH] feat: jupyter support, using cell magic

---
 README.md          | 51 ++++++++++++++++++++++++++++++++++++++++++++++
 per4m/cellmagic.py | 41 +++++++++++++++++++++++++++++++++++++
 per4m/example2.py  | 20 ++++++++++++++++++
 3 files changed, 112 insertions(+)
 create mode 100644 per4m/cellmagic.py
 create mode 100644 per4m/example2.py

diff --git a/README.md b/README.md
index 08be880..eb48a8c 100644
--- a/README.md
+++ b/README.md
@@ -18,6 +18,57 @@ Open the result.html, and identify the problem (GIL visible, possible low instru
 
 
 The dark red `S(GIL)` blocks indicate the threads/processes are in a waiting state due to the GIL, dark orange `S` is a due to other reasons (like `time.sleep(...)`). The regular pattern is due to Python switching threads after [`sys.getswitchinterval`](https://docs.python.org/3/library/sys.html#sys.getswitchinterval) (0.005 seconds)
+
+# Usage - Jupyter notebook
+
+First, load the magics
+```
+%load_ext per4m.cellmagic
+```
+
+Run a cell with the `%%giltrace` cell magic.
+```
+%%giltrace
+import threading
+import time
+import time
+
+
+def run():
+    total = 0
+    for i in range(1_000_000):
+        total += i
+    return total
+
+
+thread1 = threading.Thread(target=run)
+thread2 = threading.Thread(target=run)
+thread1.start()
+thread2.start()
+time.sleep(0.2)
+for thread in [thread1, thread2]:
+    thread.join()
+```
+Output:
+```
+Saving report to /tmp/tmp2rwf1xq3/viztracer.json ...
+Dumping trace data to json, total entries: 89, estimated json file size: 10.4KiB
+Report saved.
+
+[ perf record: Woken up 8 times to write data ]
+[ perf record: Captured and wrote 2,752 MB /tmp/tmp2rwf1xq3/perf.data (415 samples) ]
+
+Wait for perf to finish...
+Saving report to /home/maartenbreddels/github/maartenbreddels/per4m/result.html ...
+Dumping trace data to json, total entries: 167, estimated json file size: 19.6KiB
+Generating HTML report
+Report saved.
+Download result.html
+Open result.html in new tab (might not work due to security issue)
+```
+
+Click the download link to get the results.
+
 # Usage - manual
 
 ## Step 1
diff --git a/per4m/cellmagic.py b/per4m/cellmagic.py
new file mode 100644
index 0000000..22eb045
--- /dev/null
+++ b/per4m/cellmagic.py
@@ -0,0 +1,41 @@
+import os
+import tempfile
+import viztracer
+from viztracer.report_builder import ReportBuilder
+from IPython.display import HTML, display
+
+from IPython.core.magic import (cell_magic,
+                                magics_class,
+                                Magics,
+                                needs_local_scope,
+                                )
+
+
+from .giltracer import GilTracer
+
+@magics_class
+class GilTraceMagic(Magics):
+    @needs_local_scope
+    @cell_magic
+    def giltrace(self, line, cell, local_ns):
+        temp_dir = tempfile.mkdtemp()
+        perf_path = os.path.join(temp_dir, 'perf.data')
+        viz_path = os.path.join(temp_dir, 'viztracer.json')
+        gil_path = os.path.join(temp_dir, 'giltracer.json')
+        out_path = 'result.html'
+        code = self.shell.transform_cell(cell)
+        with GilTracer(perf_path, gil_path) as gt:
+            with viztracer.VizTracer(output_file=viz_path):
+                exec(code, local_ns, local_ns)
+        builder = ReportBuilder([viz_path, gil_path])
+        builder.save(output_file=out_path)
+        
+        download = HTML(f'''<a href="{out_path}" download>Download {out_path}</a>''')
+        view = HTML(f'''<a href="{out_path}" target="_blank" rel=”noopener noreferrer”>Open {out_path} in new tab</a> (might not work due to security issue)''')
+        display(download, view)
+
+def load_ipython_extension(ipython):
+    """
+    Use `%load_ext per4m.cellmagic`
+    """
+    ipython.register_magics(GilTraceMagic)
\ No newline at end of file
diff --git a/per4m/example2.py b/per4m/example2.py
new file mode 100644
index 0000000..66b2367
--- /dev/null
+++ b/per4m/example2.py
@@ -0,0 +1,20 @@
+# same as example1, but without explicit viztracer calls
+import threading
+import time
+import time
+
+
+def run():
+    total = 0
+    for i in range(1_000_000):
+        total += i
+    return total
+
+
+thread1 = threading.Thread(target=run)
+thread2 = threading.Thread(target=run)
+thread1.start()
+thread2.start()
+time.sleep(0.2)
+for thread in [thread1, thread2]:
+    thread.join()