Inverted flame graph

Added inverted flame graph which is generated with --inverted command line argument. The tree's nodes with the same parent and location are merged when we check 'Hide Import System Frames' checkbox.

Signed-off-by: Ivona Stojanovic <[email protected]>

src/memray/_ipython/flamegraph.py

@@ -90,6 +90,13 @@ def argument_parser() -> argparse.ArgumentParser:
         default=False,
     )
 
+    parser.add_argument(
+        "--inverted",
+        help="Invert flame graph",
+        action="store_true",
+        default=False,
+    )
+
     return parser
 
 
@@ -148,6 +155,7 @@ def memray_flamegraph(self, line: str, cell: str) -> None:
                 snapshot,
                 memory_records=memory_records,
                 native_traces=options.native,
+                inverted=options.inverted,
             )
         assert reporter is not None
         flamegraph_path = Path(tempdir) / "flamegraph.html"

src/memray/commands/common.py

@@ -123,11 +123,13 @@ def write_report(
         show_memory_leaks: bool,
         temporary_allocation_threshold: int,
         merge_threads: Optional[bool] = None,
+        inverted: Optional[bool] = None,
         temporal: bool = False,
     ) -> None:
         try:
             reader = FileReader(os.fspath(result_path), report_progress=True)
             merge_threads = True if merge_threads is None else merge_threads
+            inverted = False if inverted is None else inverted
 
             if reader.metadata.has_native_traces:
                 warn_if_not_enough_symbols()
@@ -143,6 +145,7 @@ def write_report(
                         memory_records=tuple(reader.get_memory_snapshots()),
                         native_traces=reader.metadata.has_native_traces,
                         high_water_mark_by_snapshot=None,
+                        inverted=inverted,
                     )
                 else:
                     recs, hwms = reader.get_temporal_high_water_mark_allocation_records(
@@ -153,6 +156,7 @@ def write_report(
                         memory_records=tuple(reader.get_memory_snapshots()),
                         native_traces=reader.metadata.has_native_traces,
                         high_water_mark_by_snapshot=hwms,
+                        inverted=inverted,
                     )
             else:
                 if show_memory_leaks:
@@ -172,6 +176,7 @@ def write_report(
                     snapshot,
                     memory_records=tuple(reader.get_memory_snapshots()),
                     native_traces=reader.metadata.has_native_traces,
+                    inverted=inverted,
                 )
         except OSError as e:
             raise MemrayCommandError(
@@ -185,6 +190,7 @@ def write_report(
                 metadata=reader.metadata,
                 show_memory_leaks=show_memory_leaks,
                 merge_threads=merge_threads,
+                inverted=inverted,
             )
 
     def prepare_parser(self, parser: argparse.ArgumentParser) -> None:
@@ -261,6 +267,9 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
         if hasattr(args, "split_threads"):
             kwargs["merge_threads"] = not args.split_threads
 
+        if hasattr(args, "inverted"):
+            kwargs["inverted"] = args.inverted
+
         self.write_report(
             result_path,
             output_file,

src/memray/commands/flamegraph.py

@@ -22,3 +22,10 @@ def prepare_parser(self, parser: argparse.ArgumentParser) -> None:
             action="store_true",
             default=False,
         )
+
+        parser.add_argument(
+            "--inverted",
+            help="Invert flame graph",
+            action="store_true",
+            default=False,
+        )

src/memray/reporters/assets/common.js

@@ -192,6 +192,32 @@ function filterFramesByFunc(root, func) {
   return _.defaults({ children: children }, root);
 }
 
+function mergeNodes(root) {
+  for (let i = 0; i < root.children.length; i++) {
+    const node = root.children[i];
+    const same_nodes = root.children.filter(
+      (n) => n.location == node.location && n != node
+    );
+
+    if (!same_nodes.length) {
+      continue;
+    }
+
+    for (let same_node of same_nodes) {
+      node.children = node.children.concat(same_node.children);
+      node.n_allocations += same_node.n_allocations;
+      node.value += same_node.value;
+    }
+
+    root.children = root.children.filter(
+      (n) => n.location != node.location || n == node
+    );
+    node.children = mergeNodes(node).children;
+  }
+
+  return root;
+}
+
 export function filterUninteresting(root) {
   return filterFramesByFunc(root, (node) => {
     return node.interesting;
@@ -204,6 +230,10 @@ export function filterImportSystem(root) {
   });
 }
 
+export function filterMergeNodesInverted(root) {
+  return mergeNodes(root);
+}
+
 /**
  * Walk the children of the specified node and sum up the total allocations and memory use.
  *

src/memray/reporters/assets/flamegraph.js

@@ -11,6 +11,7 @@ import {
   onResize,
   onInvert,
   getFlamegraph,
+  setMergeNodesInvertedFlamegraph,
 } from "./flamegraph_common";
 
 window.resizeMemoryGraph = resizeMemoryGraph;
@@ -66,6 +67,7 @@ function main() {
   hideImportSystemCheckBox.onclick = onFilterImportSystem.bind(this);
   // Enable filtering by default
   onFilterUninteresting.bind(this)();
+  setMergeNodesInvertedFlamegraph.bind(this)();
 
   document.onkeyup = (event) => {
     if (event.code == "Escape") {

src/memray/reporters/assets/flamegraph_common.js

@@ -5,6 +5,7 @@ import {
   humanFileSize,
   makeTooltipString,
   sumAllocations,
+  filterMergeNodesInverted,
 } from "./common";
 
 const FILTER_UNINTERESTING = "filter_uninteresting";
@@ -22,11 +23,18 @@ class FilteredChart {
     delete this.filters[name];
   }
 
+  setInverted(inverted) {
+    this.inverted = inverted;
+  }
+
   drawChart(data) {
     let filtered = data;
     _.forOwn(this.filters, (func) => {
       filtered = func(filtered);
     });
+    if (this.inverted) {
+      filtered = filterMergeNodesInverted(filtered);
+    }
     drawChart(filtered);
     // Merge 0 additional elements, triggering a redraw
     chart.merge([]);
@@ -98,6 +106,10 @@ export function onResize() {
   }
 }
 
+export function setMergeNodesInvertedFlamegraph() {
+  filteredChart.setInverted(inverted);
+}
+
 export function onFilterThread() {
   const thread_id = this.dataset.thread;
   if (thread_id === "-0x1") {

src/memray/reporters/flamegraph.py

@@ -75,6 +75,21 @@ def __init__(
         self.data = data
         self.memory_records = memory_records
 
+    @classmethod
+    def update_import_system_frames_for_inverted_flamegraph(
+        cls, nodes: Dict[str, Any], node_index: int
+    ) -> bool:
+        if not nodes or not nodes[node_index]:
+            return False
+
+        for child_index in nodes[node_index]["children"]:
+            result = cls.update_import_system_frames_for_inverted_flamegraph(
+                nodes, child_index
+            )
+            nodes[node_index]["import_system"] |= result
+
+        return nodes[node_index]["import_system"]
+
     @classmethod
     def _from_any_snapshot(
         cls,
@@ -83,6 +98,7 @@ def _from_any_snapshot(
         memory_records: Iterable[MemorySnapshot],
         native_traces: bool,
         temporal: bool,
+        inverted: Optional[bool] = None,
     ) -> "FlameGraphReporter":
         root: Dict[str, Any] = {
             "name": "<root>",
@@ -132,14 +148,15 @@ def _from_any_snapshot(
             )
             num_skipped_frames = 0
             is_import_system = False
-            for index, stack_frame in enumerate(reversed(stack)):
+            stack_it = enumerate((stack)) if inverted else enumerate(reversed(stack))
+            for index, stack_frame in stack_it:
                 node_key = (current_frame_id, stack_frame, thread_id)
                 if node_key not in node_index_by_key:
                     if is_cpython_internal(stack_frame):
                         num_skipped_frames += 1
                         continue
 
-                    if not is_import_system:
+                    if inverted or not is_import_system:
                         is_import_system = is_frame_from_import_system(stack_frame)
 
                     new_node_id = len(frames)
@@ -178,6 +195,11 @@ def _from_any_snapshot(
                     for interval in intervals
                 )
 
+        if inverted:
+            frames[0][
+                "import_system"
+            ] &= cls.update_import_system_frames_for_inverted_flamegraph(frames, 0)
+
         all_strings = StringRegistry()
         nodes = collections.defaultdict(list)
         for frame in frames:
@@ -213,12 +235,14 @@ def from_snapshot(
         *,
         memory_records: Iterable[MemorySnapshot],
         native_traces: bool,
+        inverted: Optional[bool] = None,
     ) -> "FlameGraphReporter":
         return cls._from_any_snapshot(
             allocations,
             memory_records=memory_records,
             native_traces=native_traces,
             temporal=False,
+            inverted=inverted,
         )
 
     @classmethod
@@ -229,12 +253,14 @@ def from_temporal_snapshot(
         memory_records: Iterable[MemorySnapshot],
         native_traces: bool,
         high_water_mark_by_snapshot: Optional[List[int]],
+        inverted: Optional[bool] = None,
     ) -> "FlameGraphReporter":
         ret = cls._from_any_snapshot(
             allocations,
             memory_records=memory_records,
             native_traces=native_traces,
             temporal=True,
+            inverted=inverted,
         )
         ret.data["high_water_mark_by_snapshot"] = high_water_mark_by_snapshot
         return ret
@@ -245,6 +271,7 @@ def render(
         metadata: Metadata,
         show_memory_leaks: bool,
         merge_threads: bool,
+        inverted: bool,
     ) -> None:
         kind = "temporal_flamegraph" if "intervals" in self.data else "flamegraph"
         html_code = render_report(
@@ -254,5 +281,6 @@ def render(
             memory_records=self.memory_records,
             show_memory_leaks=show_memory_leaks,
             merge_threads=merge_threads,
+            inverted=inverted,
         )
         print(html_code, file=outfile)

src/memray/reporters/templates/__init__.py

@@ -28,7 +28,9 @@ def include_file(name: str) -> Markup:
     return env
 
 
-def get_report_title(*, kind: str, show_memory_leaks: bool) -> str:
+def get_report_title(*, kind: str, show_memory_leaks: bool, inverted: bool) -> str:
+    if inverted:
+        kind = f"inverted {kind}"
     if show_memory_leaks:
         return f"{kind} report (memory leaks)"
     return f"{kind} report"
@@ -42,12 +44,17 @@ def render_report(
     memory_records: Iterable[MemorySnapshot],
     show_memory_leaks: bool,
     merge_threads: bool,
+    inverted: bool,
 ) -> str:
     env = get_render_environment()
     template = env.get_template(kind + ".html")
 
     pretty_kind = kind.replace("_", " ")
-    title = get_report_title(kind=pretty_kind, show_memory_leaks=show_memory_leaks)
+    title = get_report_title(
+        kind=pretty_kind,
+        show_memory_leaks=show_memory_leaks,
+        inverted=inverted,
+    )
     return template.render(
         kind=pretty_kind,
         title=title,
@@ -56,4 +63,5 @@ def render_report(
         memory_records=memory_records,
         show_memory_leaks=show_memory_leaks,
         merge_threads=merge_threads,
+        inverted=inverted,
     )

src/memray/reporters/templates/base.html

@@ -141,6 +141,7 @@ <h5 class="modal-title" id="helpModalLabel">How to interpret {{ kind }} reports<
     var data = null;
     const merge_threads = {{ merge_threads|tojson }};
     const memory_records = {{ memory_records|tojson }};
+    const inverted = {{ inverted|tojson }};
   </script>
   {% endblock scripts %}
 </body>